From 75f5399d6f8f50b2920666d56b4e7ebf9e1ddf3c Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Mon, 6 Jul 2026 15:06:33 +0800 Subject: [PATCH 1/8] refactor: clarify module boundaries for property types and export IO --- extension/gds/include/impl/leiden_impl.h | 2 +- extension/gds/include/impl/louvain_impl.h | 2 +- extension/gds/include/utils/parallel_utils.h | 2 +- .../parquet/include/parquet/arrow_reader.h | 18 +- .../parquet/include/parquet_export_function.h | 3 +- .../parquet/include/parquet_read_function.h | 2 +- extension/parquet/src/arrow_reader.cc | 28 +- .../parquet/src/parquet_export_function.cc | 693 ++++++++++-------- .../parquet/tests/parquet_export_test_utils.h | 252 +++++++ extension/parquet/tests/parquet_test.cc | 73 +- .../fastest_lib/src/DataStructure/graph.h | 2 +- include/neug/common/columns/edge_columns.h | 3 +- include/neug/common/columns/value_columns.h | 2 +- .../property => common/types}/default_value.h | 0 include/neug/common/types/graph_types.h | 2 +- include/neug/common/types/i_context_column.h | 2 +- .../types.h => common/types/property_types.h} | 0 .../function/export/json_export_function.h | 39 +- .../function/import/csv_read_function.h | 2 +- .../function/import/json_read_function.h | 4 +- include/neug/execution/common/context.h | 9 + .../common/operators/retrieve/sink.h | 7 + .../execute/ops/batch/batch_update_utils.h | 2 +- .../execute/ops/retrieve/scan_utils.h | 2 +- .../execution/expression/special_predicates.h | 2 +- include/neug/main/neug_db.h | 2 +- include/neug/storages/README.md | 4 +- include/neug/storages/allocators.h | 2 +- include/neug/storages/container/i_container.h | 2 +- include/neug/storages/csr/csr_base.h | 2 +- include/neug/storages/csr/csr_view.h | 4 +- include/neug/storages/csr/immutable_csr.h | 2 +- include/neug/storages/csr/mutable_csr.h | 2 +- include/neug/storages/csr/nbr.h | 6 +- include/neug/storages/graph/edge_table.h | 6 +- include/neug/storages/graph/graph_interface.h | 2 +- include/neug/storages/graph/graph_view.h | 2 +- .../neug/storages/graph/operation_params.h | 2 +- include/neug/storages/graph/property_graph.h | 2 +- .../storages/graph/property_graph_cow_state.h | 2 +- include/neug/storages/graph/schema.h | 4 +- include/neug/storages/graph/vertex_table.h | 4 +- .../neug/storages/graph/vertex_timestamp.h | 2 +- include/neug/{utils => storages}/id_indexer.h | 4 +- include/neug/{utils => storages}/indexers.h | 4 +- .../loader/abstract_property_graph_loader.h | 2 +- .../neug/storages/loader/i_fragment_loader.h | 4 +- include/neug/storages/loader/loading_config.h | 2 +- include/neug/storages/module/type_name.h | 4 +- .../property/array_column.h | 2 +- .../{utils => storages}/property/column.h | 2 +- .../neug/{utils => storages}/property/table.h | 4 +- .../neug/transaction/compact_transaction.h | 2 +- include/neug/transaction/insert_transaction.h | 2 +- include/neug/transaction/read_transaction.h | 6 +- include/neug/transaction/transaction_utils.h | 2 +- include/neug/transaction/update_transaction.h | 4 +- include/neug/transaction/wal/wal.h | 2 +- include/neug/transaction/wal/wal_builder.h | 2 +- .../utils/io/read/common/type_converter.h | 2 +- .../neug/utils/io/read/csv/csv_read_config.h | 2 +- include/neug/utils/io/read/csv/csv_reader.h | 18 +- .../utils/io/read/json/json_read_config.h | 2 +- include/neug/utils/io/read/json/json_reader.h | 18 +- include/neug/utils/io/write/writer.h | 70 +- include/neug/utils/pb_utils.h | 2 +- include/neug/utils/serialization/in_archive.h | 2 +- .../neug/utils/serialization/out_archive.h | 2 +- include/neug/utils/service_utils.h | 2 +- include/neug/utils/string_utils.h | 2 +- .../types}/default_value.cc | 2 +- src/common/types/graph_types.cc | 2 +- .../types/property_types.cc} | 4 +- src/compiler/function/csv_export_function.cpp | 15 +- src/compiler/function/json_export_function.cc | 468 +++++------- src/execution/common/context.cc | 22 + .../common/operators/retrieve/join.cc | 2 +- .../operators/retrieve/path_expand_impl.cc | 2 +- .../common/operators/retrieve/sink.cc | 157 +++- src/execution/execute/ops/retrieve/edge.cc | 2 +- .../execute/ops/retrieve/group_by.cc | 2 +- src/execution/execute/ops/retrieve/scan.cc | 2 +- src/execution/execute/ops/retrieve/select.cc | 2 +- src/execution/execute/ops/retrieve/tc_fuse.cc | 2 +- src/execution/execute/ops/retrieve/vertex.cc | 2 +- src/execution/utils/pb_parse_utils.cc | 2 +- src/main/session.cc | 2 +- src/storages/CMakeLists.txt | 3 +- src/storages/csr/csr_view_utils.cc | 2 +- src/storages/csr/immutable_csr.cc | 2 +- src/storages/csr/mutable_csr.cc | 2 +- src/storages/graph/edge_table.cc | 2 +- src/storages/graph/property_graph.cc | 6 +- src/storages/graph/schema.cc | 4 +- src/storages/loader/loader_utils.cc | 2 +- src/storages/property/CMakeLists.txt | 6 + .../property/array_column.cc | 4 +- src/{utils => storages}/property/column.cc | 10 +- src/{utils => storages}/property/table.cc | 4 +- src/transaction/insert_transaction.cc | 2 +- src/transaction/update_transaction.cc | 8 +- src/utils/bolt_utils.cc | 2 +- src/utils/io/read/csv/csv_reader.cc | 29 +- src/utils/io/read/json/json_reader.cc | 29 +- src/utils/io/write/writer.cc | 341 +++------ src/utils/pb_utils.cc | 4 +- src/utils/yaml_utils.cc | 2 +- tests/storage/alter_property_test.cc | 2 +- tests/storage/test_vertex_table.cc | 2 +- tests/unittest/logical_delete_test.cc | 2 +- tests/unittest/schema_test.cc | 2 +- tests/unittest/test_column.cc | 4 +- tests/unittest/test_indexer.cc | 2 +- tests/unittest/utils.h | 8 +- tests/utils/json_test.cc | 2 +- tests/utils/test_reader.cc | 34 +- tests/utils/test_table.cc | 4 +- tests/utils/test_types.cc | 4 +- tools/nodejs_bind/src/node_query_request.cc | 2 +- tools/nodejs_bind/src/node_query_result.cc | 2 +- tools/python_bind/src/py_query_request.cc | 2 +- tools/python_bind/src/py_query_result.cc | 2 +- 122 files changed, 1405 insertions(+), 1198 deletions(-) create mode 100644 extension/parquet/tests/parquet_export_test_utils.h rename include/neug/{utils/property => common/types}/default_value.h (100%) rename include/neug/{utils/property/types.h => common/types/property_types.h} (100%) rename include/neug/{utils => storages}/id_indexer.h (99%) rename include/neug/{utils => storages}/indexers.h (89%) rename include/neug/{utils => storages}/property/array_column.h (98%) rename include/neug/{utils => storages}/property/column.h (99%) rename include/neug/{utils => storages}/property/table.h (97%) rename src/{utils/property => common/types}/default_value.cc (97%) rename src/{utils/property/types.cc => common/types/property_types.cc} (99%) create mode 100644 src/storages/property/CMakeLists.txt rename src/{utils => storages}/property/array_column.cc (98%) rename src/{utils => storages}/property/column.cc (94%) rename src/{utils => storages}/property/table.cc (99%) diff --git a/extension/gds/include/impl/leiden_impl.h b/extension/gds/include/impl/leiden_impl.h index 11cf41f31..5dca141ed 100644 --- a/extension/gds/include/impl/leiden_impl.h +++ b/extension/gds/include/impl/leiden_impl.h @@ -19,9 +19,9 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" namespace neug { namespace gds { diff --git a/extension/gds/include/impl/louvain_impl.h b/extension/gds/include/impl/louvain_impl.h index 9df0457fd..fb5925769 100644 --- a/extension/gds/include/impl/louvain_impl.h +++ b/extension/gds/include/impl/louvain_impl.h @@ -19,9 +19,9 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" namespace neug { namespace gds { diff --git a/extension/gds/include/utils/parallel_utils.h b/extension/gds/include/utils/parallel_utils.h index 466bf7261..9f26248d5 100644 --- a/extension/gds/include/utils/parallel_utils.h +++ b/extension/gds/include/utils/parallel_utils.h @@ -21,8 +21,8 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/graph/vertex_table.h" -#include "neug/utils/property/types.h" namespace neug { namespace gds { diff --git a/extension/parquet/include/parquet/arrow_reader.h b/extension/parquet/include/parquet/arrow_reader.h index 45666e471..70676979c 100644 --- a/extension/parquet/include/parquet/arrow_reader.h +++ b/extension/parquet/include/parquet/arrow_reader.h @@ -15,12 +15,14 @@ #pragma once #include +#include #include #include #include #include +#include "neug/common/types/data_chunk.h" #include "neug/utils/io/reader.h" #include "parquet/arrow_options.h" @@ -47,8 +49,8 @@ class Reader { fileSystem(std::move(fileSystem)) {} virtual ~Reader() = default; - virtual void read(std::shared_ptr localState, - execution::Context& ctx) = 0; + virtual std::vector read( + std::shared_ptr localState) = 0; protected: std::shared_ptr sharedState; @@ -72,18 +74,18 @@ class ArrowReader : public Reader { datasetBuilder(std::move(datasetBuilder)) {} ~ArrowReader() override = default; - void read(std::shared_ptr localState, - execution::Context& ctx) override; + std::vector read( + std::shared_ptr localState) 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); + std::vector full_read( + std::shared_ptr scanner); + std::vector batch_read( + std::shared_ptr scanner); std::unique_ptr optionsBuilder; std::shared_ptr datasetBuilder; diff --git a/extension/parquet/include/parquet_export_function.h b/extension/parquet/include/parquet_export_function.h index b5b22e476..3d243d1ef 100644 --- a/extension/parquet/include/parquet_export_function.h +++ b/extension/parquet/include/parquet_export_function.h @@ -38,7 +38,8 @@ class ArrowParquetExportWriter : public QueryExportWriter { fileSystem_(std::move(fileSystem)) {} ~ArrowParquetExportWriter() override = default; - neug::Status writeTable(const QueryResponse* table) override; + neug::Status write(const DataChunk& chunk, + const std::vector& source_types = {}) override; private: std::shared_ptr fileSystem_; diff --git a/extension/parquet/include/parquet_read_function.h b/extension/parquet/include/parquet_read_function.h index 43177547b..573a1b140 100644 --- a/extension/parquet/include/parquet_read_function.h +++ b/extension/parquet/include/parquet_read_function.h @@ -66,7 +66,7 @@ struct ParquetReadFunction { execution::Context ctx; auto localState = std::make_shared(); - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); return ctx; } diff --git a/extension/parquet/src/arrow_reader.cc b/extension/parquet/src/arrow_reader.cc index f9b313651..3b4c18752 100644 --- a/extension/parquet/src/arrow_reader.cc +++ b/extension/parquet/src/arrow_reader.cc @@ -24,7 +24,6 @@ #include "parquet/record_batch_supplier.h" #include "neug/compiler/common/assert.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" @@ -32,8 +31,8 @@ namespace neug { namespace reader { -void ArrowReader::read(std::shared_ptr localState, - execution::Context& ctx) { +std::vector ArrowReader::read( + std::shared_ptr localState) { if (!sharedState) { THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); } @@ -49,10 +48,9 @@ void ArrowReader::read(std::shared_ptr localState, const auto& fileSchema = sharedState->schema.file; ReadOptions options; if (options.batch_read.get(fileSchema.options)) { - batch_read(scanner, ctx); - } else { - full_read(scanner, ctx); + return batch_read(scanner); } + return full_read(scanner); } std::shared_ptr ArrowReader::createScanner( @@ -136,8 +134,8 @@ std::shared_ptr ArrowReader::createScanner( return scanner_result.ValueOrDie(); } -void ArrowReader::full_read(std::shared_ptr scanner, - execution::Context& output) { +std::vector ArrowReader::full_read( + std::shared_ptr scanner) { if (!sharedState) { THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); } @@ -162,17 +160,18 @@ void ArrowReader::full_read(std::shared_ptr scanner, ", 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(std::move(chunk)); + std::vector output; + output.push_back(std::move(chunk)); + return output; } -void ArrowReader::batch_read(std::shared_ptr scanner, - execution::Context& output) { +std::vector ArrowReader::batch_read( + std::shared_ptr scanner) { if (!sharedState) { THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); } @@ -203,10 +202,11 @@ void ArrowReader::batch_read(std::shared_ptr scanner, auto batch_supplier = std::make_shared(batch_reader, row_num); - output.clear(); + std::vector output; while (auto chunk = batch_supplier->GetNextChunk()) { - output.append_chunk(std::move(*chunk)); + output.push_back(std::move(*chunk)); } + return output; } arrow::Result> ArrowReader::inferSchema() { diff --git a/extension/parquet/src/parquet_export_function.cc b/extension/parquet/src/parquet_export_function.cc index ce770c789..8967b314d 100644 --- a/extension/parquet/src/parquet_export_function.cc +++ b/extension/parquet/src/parquet_export_function.cc @@ -23,22 +23,25 @@ #include #include #include -#include -#include -#include #include - +#include + +#include "neug/common/columns/list_columns.h" +#include "neug/common/columns/struct_columns.h" +#include "neug/common/columns/value_columns.h" +#include "neug/common/types/array_columns.h" +#include "neug/common/types/data_chunk.h" +#include "neug/common/types/property_types.h" +#include "neug/common/types/value.h" #include "neug/compiler/main/metadata_registry.h" +#include "neug/execution/common/operators/retrieve/sink.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/io/write/writer.h" -#include "neug/utils/property/types.h" #include "parquet/arrow_fs_resolver.h" #include "parquet_options.h" namespace neug { namespace writer { -// Parse writer options and build WriterProperties static std::shared_ptr<::parquet::WriterProperties> buildWriterProperties( const common::case_insensitive_map_t& options) { reader::ParquetExportOptions export_options; @@ -95,362 +98,432 @@ static std::shared_ptr<::parquet::WriterProperties> buildWriterProperties( return builder.build(); } -// Infer Arrow type from protobuf Array structure -static std::shared_ptr inferArrowTypeFromArray( - const Array& proto_array) { - if (proto_array.has_int32_array()) { +static std::shared_ptr inferArrowTypeFromDataType( + const DataType& type) { + switch (type.id()) { + case DataTypeId::kBoolean: + return arrow::boolean(); + case DataTypeId::kInt32: return arrow::int32(); - } else if (proto_array.has_int64_array()) { + case DataTypeId::kInt64: return arrow::int64(); - } else if (proto_array.has_uint32_array()) { + case DataTypeId::kUInt32: return arrow::uint32(); - } else if (proto_array.has_uint64_array()) { + case DataTypeId::kUInt64: return arrow::uint64(); - } else if (proto_array.has_float_array()) { + case DataTypeId::kFloat: return arrow::float32(); - } else if (proto_array.has_double_array()) { + case DataTypeId::kDouble: return arrow::float64(); - } else if (proto_array.has_bool_array()) { - return arrow::boolean(); - } else if (proto_array.has_string_array()) { + case DataTypeId::kVarchar: return arrow::large_utf8(); - } else if (proto_array.has_date_array()) { + case DataTypeId::kDate: return arrow::date64(); - } else if (proto_array.has_timestamp_array()) { + case DataTypeId::kTimestampMs: return arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"); - } else if (proto_array.has_list_array()) { - // Recursively infer element type - const auto& list_arr = proto_array.list_array(); - if (list_arr.has_elements()) { - auto element_type = inferArrowTypeFromArray(list_arr.elements()); - return arrow::list(element_type); - } - return arrow::list(arrow::large_utf8()); // Default to string list - } else if (proto_array.has_struct_array()) { - // Struct type: infer types from each field - const auto& struct_arr = proto_array.struct_array(); + case DataTypeId::kInterval: + return arrow::large_utf8(); + case DataTypeId::kList: + return arrow::list( + inferArrowTypeFromDataType(ListType::GetChildType(type))); + case DataTypeId::kArray: + return arrow::list( + inferArrowTypeFromDataType(ArrayType::GetChildType(type))); + case DataTypeId::kStruct: { std::vector> fields; - for (int i = 0; i < struct_arr.fields_size(); ++i) { - auto field_type = inferArrowTypeFromArray(struct_arr.fields(i)); - fields.push_back(arrow::field("field_" + std::to_string(i), field_type)); + const auto& child_types = StructType::GetChildTypes(type); + const auto& field_names = StructType::GetFieldNames(type); + for (size_t i = 0; i < child_types.size(); ++i) { + const auto& name = i < field_names.size() + ? field_names[i] + : ("field_" + std::to_string(i)); + fields.push_back( + arrow::field(name, inferArrowTypeFromDataType(child_types[i]))); } return arrow::struct_(fields); - } else if (proto_array.has_vertex_array() || proto_array.has_edge_array() || - proto_array.has_path_array()) { - // Vertex/Edge/Path are exported as JSON strings (not StructArray) - return arrow::large_utf8(); - } else if (proto_array.has_interval_array()) { - // Interval type: convert to string for Parquet compatibility - return arrow::large_utf8(); - } else { - LOG(WARNING) << "Unknown protobuf array type, defaulting to large_utf8"; + } + default: + LOG(WARNING) << "Unknown DataType for Parquet export, defaulting to " + "large_utf8: " + << type.ToString(); return arrow::large_utf8(); } } -// Macro for primitive array conversion (proto-type based dispatch) -#define TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(PROTO_FIELD, BUILDER_TYPE, \ - VALUES_FIELD) \ - { \ - auto& arr = proto_array.PROTO_FIELD(); \ - BUILDER_TYPE builder(pool); \ - for (int i = 0; i < arr.values_size(); ++i) { \ - if (writer::StringFormatBuffer::validateProtoValue(arr.validity(), i)) { \ - auto status = builder.Append(arr.VALUES_FIELD(i)); \ - if (!status.ok()) { \ - THROW_RUNTIME_ERROR("Failed to append value: " + status.ToString()); \ - } \ - } else { \ - auto status = builder.AppendNull(); \ - if (!status.ok()) { \ - THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); \ - } \ - } \ - } \ - std::shared_ptr result; \ - auto status = builder.Finish(&result); \ - if (!status.ok()) { \ - THROW_RUNTIME_ERROR("Failed to finish array: " + status.ToString()); \ - } \ - return result; \ +template +static std::shared_ptr buildPrimitiveArray( + const std::shared_ptr& col) { + auto typed = std::dynamic_pointer_cast>(col); + if (!typed) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected ValueColumn for type " + + col->column_info()); } - -// Convert protobuf Array to Arrow Array -static std::shared_ptr protoArrayToArrowArray( - const Array& proto_array, - const std::shared_ptr& arrow_type, int row_count) { arrow::MemoryPool* pool = arrow::default_memory_pool(); + Builder builder(pool); + for (size_t i = 0; i < typed->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + auto status = builder.AppendNull(); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); + } + } else { + auto status = builder.Append(typed->get_value(i)); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append value: " + status.ToString()); + } + } + } + std::shared_ptr result; + auto status = builder.Finish(&result); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to finish array: " + status.ToString()); + } + return result; +} - // First, dispatch based on proto array type - if (proto_array.has_int32_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(int32_array, arrow::Int32Builder, - values) - } else if (proto_array.has_int64_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(int64_array, arrow::Int64Builder, - values) - } else if (proto_array.has_uint32_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(uint32_array, arrow::UInt32Builder, - values) - } else if (proto_array.has_uint64_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(uint64_array, arrow::UInt64Builder, - values) - } else if (proto_array.has_float_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(float_array, arrow::FloatBuilder, - values) - } else if (proto_array.has_double_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(double_array, arrow::DoubleBuilder, - values) - } else if (proto_array.has_bool_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(bool_array, arrow::BooleanBuilder, - values) - } else if (proto_array.has_string_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(string_array, arrow::LargeStringBuilder, - values) - } else if (proto_array.has_date_array()) { - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(date_array, arrow::Date64Builder, - values) - } else if (proto_array.has_interval_array()) { - // Interval: convert to string for Parquet compatibility - TYPED_PRIMITIVE_ARRAY_TO_ARROW_IMPL(interval_array, - arrow::LargeStringBuilder, values) - } else if (proto_array.has_timestamp_array()) { - auto& arr = proto_array.timestamp_array(); - arrow::TimestampBuilder builder( - arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), pool); - for (int i = 0; i < arr.values_size(); ++i) { - if (writer::StringFormatBuffer::validateProtoValue(arr.validity(), i)) { - auto status = builder.Append(arr.values(i)); +static std::shared_ptr boolVectorToArrowBitmap( + const vector_t& validity, size_t length) { + if (validity.empty()) { + return nullptr; + } + auto buffer_result = + arrow::AllocateBuffer(static_cast((length + 7) / 8)); + if (!buffer_result.ok()) { + THROW_RUNTIME_ERROR("Failed to allocate validity buffer: " + + buffer_result.status().ToString()); + } + auto buffer = std::move(buffer_result.ValueOrDie()); + memset(buffer->mutable_data(), 0, static_cast(buffer->size())); + for (size_t i = 0; i < length; ++i) { + if (validity[i]) { + buffer->mutable_data()[i / 8] |= static_cast(1U << (i % 8)); + } + } + return buffer; +} + +static std::shared_ptr contextColumnToArrowArray( + const std::shared_ptr& col, + const std::shared_ptr& arrow_type) { + if (!col || col->size() == 0) { + return arrow::MakeArrayOfNull(arrow_type, 0).ValueOrDie(); + } + + switch (col->elem_type().id()) { + case DataTypeId::kBoolean: + return buildPrimitiveArray(col); + case DataTypeId::kInt32: + return buildPrimitiveArray(col); + case DataTypeId::kInt64: + return buildPrimitiveArray(col); + case DataTypeId::kUInt32: + return buildPrimitiveArray(col); + case DataTypeId::kUInt64: + return buildPrimitiveArray(col); + case DataTypeId::kFloat: + return buildPrimitiveArray(col); + case DataTypeId::kDouble: + return buildPrimitiveArray(col); + case DataTypeId::kVarchar: + return buildPrimitiveArray(col); + case DataTypeId::kDate: { + auto typed = std::dynamic_pointer_cast>(col); + if (!typed) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected ValueColumn"); + } + arrow::Date64Builder builder(arrow::default_memory_pool()); + for (size_t i = 0; i < typed->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + auto status = builder.AppendNull(); if (!status.ok()) { - THROW_RUNTIME_ERROR("Failed to append timestamp value: " + - status.ToString()); + THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); } } else { + auto status = builder.Append(typed->get_value(i).to_timestamp()); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append date: " + status.ToString()); + } + } + } + std::shared_ptr result; + auto finish_status = builder.Finish(&result); + if (!finish_status.ok()) { + THROW_RUNTIME_ERROR("Failed to finish date array: " + + finish_status.ToString()); + } + return result; + } + case DataTypeId::kTimestampMs: { + auto typed = std::dynamic_pointer_cast>(col); + if (!typed) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected ValueColumn"); + } + arrow::TimestampBuilder builder( + arrow::timestamp(arrow::TimeUnit::MICRO, "UTC"), + arrow::default_memory_pool()); + for (size_t i = 0; i < typed->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { auto status = builder.AppendNull(); if (!status.ok()) { THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); } + } else { + auto status = builder.Append(typed->get_value(i).milli_second * 1000); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append timestamp: " + + status.ToString()); + } } } std::shared_ptr result; - auto status = builder.Finish(&result); - if (!status.ok()) { + auto finish_status = builder.Finish(&result); + if (!finish_status.ok()) { THROW_RUNTIME_ERROR("Failed to finish timestamp array: " + - status.ToString()); + finish_status.ToString()); } return result; - } else if (proto_array.has_list_array()) { - // Handle List type - build native Arrow ListArray - const auto& list_arr = proto_array.list_array(); + } + case DataTypeId::kInterval: { + auto typed = std::dynamic_pointer_cast>(col); + if (!typed) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected ValueColumn"); + } + arrow::LargeStringBuilder builder(arrow::default_memory_pool()); + for (size_t i = 0; i < typed->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + auto status = builder.AppendNull(); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); + } + } else { + auto status = builder.Append(typed->get_value(i).to_string()); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append interval: " + + status.ToString()); + } + } + } + std::shared_ptr result; + auto finish_status = builder.Finish(&result); + if (!finish_status.ok()) { + THROW_RUNTIME_ERROR("Failed to finish interval array: " + + finish_status.ToString()); + } + return result; + } + case DataTypeId::kList: { + auto list_col = std::dynamic_pointer_cast(col); + if (!list_col) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected ListColumn"); + } + auto reordered_list = + std::dynamic_pointer_cast(list_col->reorder()); auto list_type = std::static_pointer_cast(arrow_type); - auto element_type = list_type->value_type(); - - // Recursively convert all elements - auto elements_array = - protoArrayToArrowArray(list_arr.elements(), element_type, 0); - - // Build offsets buffer - int64_t num_rows = list_arr.offsets_size() - 1; - - // Create offsets buffer - copy data to avoid dangling pointer when protobuf - // is destroyed - int64_t offsets_byte_size = list_arr.offsets_size() * sizeof(int32_t); - auto offsets_buffer_result = arrow::AllocateBuffer(offsets_byte_size); - if (!offsets_buffer_result.ok()) { - THROW_RUNTIME_ERROR("Failed to allocate offsets buffer: " + - offsets_buffer_result.status().ToString()); + auto elements_array = contextColumnToArrowArray( + reordered_list->data_column(), list_type->value_type()); + + arrow::Int32Builder offsets_builder(arrow::default_memory_pool()); + const auto& items = reordered_list->items(); + for (const auto& item : items) { + auto status = offsets_builder.Append(static_cast(item.offset)); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append list offset: " + + status.ToString()); + } + } + if (!items.empty()) { + const auto& last = items.back(); + auto status = offsets_builder.Append( + static_cast(last.offset + last.length)); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append list offset: " + + status.ToString()); + } + } else { + auto status = offsets_builder.Append(0); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append list offset: " + + status.ToString()); + } + } + std::shared_ptr offsets_array; + auto finish_status = offsets_builder.Finish(&offsets_array); + if (!finish_status.ok()) { + THROW_RUNTIME_ERROR("Failed to finish list offsets: " + + finish_status.ToString()); + } + return std::make_shared( + list_type, static_cast(items.size()), + offsets_array->data()->buffers[1], elements_array, nullptr); + } + case DataTypeId::kArray: { + auto array_col = std::dynamic_pointer_cast(col); + if (!array_col) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected ContextArrayColumn"); + } + auto list_type = std::static_pointer_cast(arrow_type); + auto elements_array = contextColumnToArrowArray(array_col->data_column(), + list_type->value_type()); + const auto array_size = array_col->array_size(); + const auto num_rows = array_col->size(); + + arrow::Int32Builder offsets_builder(arrow::default_memory_pool()); + for (size_t i = 0; i < num_rows; ++i) { + auto status = + offsets_builder.Append(static_cast(i * array_size)); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append array offset: " + + status.ToString()); + } + } + auto last_status = + offsets_builder.Append(static_cast(num_rows * array_size)); + if (!last_status.ok()) { + THROW_RUNTIME_ERROR("Failed to append array offset: " + + last_status.ToString()); + } + std::shared_ptr offsets_array; + auto finish_status = offsets_builder.Finish(&offsets_array); + if (!finish_status.ok()) { + THROW_RUNTIME_ERROR("Failed to finish array offsets: " + + finish_status.ToString()); + } + return std::make_shared( + list_type, static_cast(num_rows), + offsets_array->data()->buffers[1], elements_array, nullptr); + } + case DataTypeId::kStruct: { + auto struct_col = std::dynamic_pointer_cast(col); + if (!struct_col) { + THROW_INVALID_ARGUMENT_EXCEPTION("Expected StructColumn"); } - std::shared_ptr offsets_buffer = - std::move(offsets_buffer_result.ValueOrDie()); - memcpy(offsets_buffer->mutable_data(), list_arr.offsets().data(), - offsets_byte_size); - - // If all values are valid (no nulls), we can safely omit validity buffer - std::shared_ptr validity_buffer = nullptr; - - // Create ListArray directly - auto list_array = std::make_shared( - arrow_type, num_rows, offsets_buffer, elements_array, validity_buffer); - - return list_array; - } else if (proto_array.has_struct_array()) { - // Handle StructArray - const auto& struct_arr = proto_array.struct_array(); auto struct_type = std::static_pointer_cast(arrow_type); - - // Recursively convert each field std::vector> field_arrays; - for (int i = 0; i < struct_arr.fields_size(); ++i) { - auto field_type = struct_type->field(i)->type(); - auto field_array = - protoArrayToArrowArray(struct_arr.fields(i), field_type, row_count); - field_arrays.push_back(field_array); + const auto& children = struct_col->children(); + for (size_t i = 0; i < children.size(); ++i) { + field_arrays.push_back(contextColumnToArrowArray( + children[i], struct_type->field(static_cast(i))->type())); } - - // Build validity buffer - copy data to avoid dangling pointer when protobuf - // is destroyed - auto null_bitmap = struct_arr.validity(); - std::shared_ptr validity_buffer; - if (!null_bitmap.empty()) { - auto buffer_result = arrow::AllocateBuffer(null_bitmap.size()); - if (!buffer_result.ok()) { - THROW_RUNTIME_ERROR("Failed to allocate validity buffer: " + - buffer_result.status().ToString()); - } - validity_buffer = std::move(buffer_result.ValueOrDie()); - memcpy(validity_buffer->mutable_data(), null_bitmap.data(), - null_bitmap.size()); + std::shared_ptr validity_buffer = nullptr; + if (struct_col->is_optional()) { + validity_buffer = + boolVectorToArrowBitmap(struct_col->validity_bitmap(), col->size()); } - - // Create StructArray - num_rows should match the field arrays' length - int64_t num_rows = field_arrays.empty() ? 0 : field_arrays[0]->length(); - - auto struct_array = std::make_shared( - struct_type, num_rows, field_arrays, validity_buffer); - - return struct_array; - } else if (proto_array.has_vertex_array() || proto_array.has_edge_array() || - proto_array.has_path_array()) { - // Vertex/Edge/Path are stored as JSON strings in protobuf. - // We export them directly as JSON strings (not StructArray). - // - // NOTE: Why not convert to StructArray? - // Parquet StructArray has columnar advantages (compression, predicate - // pushdown, etc.), but requires ALL rows to have the SAME schema. - // Mixed-type vertices/edges have different properties, causing schema - // conflicts and sparse data. - - // Extract JSON strings from the appropriate field - arrow::LargeStringBuilder builder(pool); - - auto append_json_strings = [&](const auto& arr) { - for (int i = 0; i < arr.values_size(); ++i) { - if (writer::StringFormatBuffer::validateProtoValue(arr.validity(), i)) { - auto status = builder.Append(arr.values(i)); - if (!status.ok()) { - THROW_RUNTIME_ERROR("Failed to append JSON string: " + - status.ToString()); - } - } else { - auto status = builder.AppendNull(); - if (!status.ok()) { - THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); - } + return std::make_shared( + struct_type, static_cast(col->size()), field_arrays, + validity_buffer); + } + default: { + arrow::LargeStringBuilder builder(arrow::default_memory_pool()); + for (size_t i = 0; i < col->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + auto status = builder.AppendNull(); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append null: " + status.ToString()); + } + } else { + auto status = builder.Append(col->get_elem(i).to_string()); + if (!status.ok()) { + THROW_RUNTIME_ERROR("Failed to append string: " + status.ToString()); } } - }; - - if (proto_array.has_vertex_array()) { - append_json_strings(proto_array.vertex_array()); - } else if (proto_array.has_edge_array()) { - append_json_strings(proto_array.edge_array()); - } else { - append_json_strings(proto_array.path_array()); } - std::shared_ptr result; - auto status = builder.Finish(&result); - if (!status.ok()) { - THROW_RUNTIME_ERROR("Failed to finish JSON string array: " + - status.ToString()); + auto finish_status = builder.Finish(&result); + if (!finish_status.ok()) { + THROW_RUNTIME_ERROR("Failed to finish string array: " + + finish_status.ToString()); } return result; - } else { - THROW_INVALID_ARGUMENT_EXCEPTION( - "Unsupported protobuf array type for conversion"); } + } +} + +static neug::Status writeArrowTable( + const std::shared_ptr& arrow_schema, + const std::vector>& arrays, + const reader::FileSchema& schema, + const std::shared_ptr& file_system) { + auto result = file_system->OpenOutputStream(schema.paths[0]); + if (!result.ok()) { + return neug::Status( + neug::StatusCode::ERR_IO_ERROR, + "Failed to open output file: " + result.status().ToString()); + } + auto outfile = result.ValueOrDie(); + + auto properties = buildWriterProperties(schema.options); + auto writer_result = ::parquet::arrow::FileWriter::Open( + *arrow_schema, arrow::default_memory_pool(), outfile, properties); + if (!writer_result.ok()) { + return neug::Status(neug::StatusCode::ERR_IO_ERROR, + "Failed to create Parquet writer: " + + writer_result.status().ToString()); + } + auto writer = std::move(writer_result.ValueOrDie()); + + auto arrow_table = arrow::Table::Make(arrow_schema, arrays); + auto write_status = writer->WriteTable(*arrow_table, arrow_table->num_rows()); + if (!write_status.ok()) { + return neug::Status( + neug::StatusCode::ERR_IO_ERROR, + "Failed to write Parquet table: " + write_status.ToString()); + } + + auto close_status = writer->Close(); + if (!close_status.ok()) { + return neug::Status( + neug::StatusCode::ERR_IO_ERROR, + "Failed to close Parquet writer: " + close_status.ToString()); + } + + auto outfile_close_status = outfile->Close(); + if (!outfile_close_status.ok()) { + return neug::Status( + neug::StatusCode::ERR_IO_ERROR, + "Failed to close output stream: " + outfile_close_status.ToString()); + } + return neug::Status::OK(); } -neug::Status ArrowParquetExportWriter::writeTable(const QueryResponse* table) { - if (!table || table->row_count() == 0) { +neug::Status ArrowParquetExportWriter::write( + const DataChunk& chunk, const std::vector& /*source_types*/) { + if (chunk.row_num() == 0) { return neug::Status::OK(); } + if (!entry_schema_) { + return neug::Status(neug::StatusCode::ERR_INVALID_ARGUMENT, + "entry_schema is null"); + } try { - // 1. Create Arrow schema from QueryResponse (infer types from protobuf - // arrays) std::vector> fields; - int num_columns = table->arrays_size(); - - for (int i = 0; i < num_columns; ++i) { - // Get column name from QueryResponse schema or entry_schema_ - std::string column_name; - if (i < table->schema().name_size()) { - column_name = table->schema().name(i); - } else if (entry_schema_ && - i < static_cast(entry_schema_->columnNames.size())) { - column_name = entry_schema_->columnNames[i]; - } else { - column_name = "col_" + std::to_string(i); + fields.reserve(chunk.col_num()); + for (size_t i = 0; i < chunk.col_num(); ++i) { + if (chunk.columns[i] == nullptr) { + continue; } - - // Infer Arrow type from protobuf array structure - const auto& proto_array = table->arrays(i); - auto arrow_type = inferArrowTypeFromArray(proto_array); - - fields.push_back(arrow::field(column_name, arrow_type)); + std::string column_name = i < entry_schema_->columnNames.size() + ? entry_schema_->columnNames[i] + : ("col_" + std::to_string(i)); + fields.push_back(arrow::field( + column_name, + inferArrowTypeFromDataType(chunk.columns[i]->elem_type()))); } auto arrow_schema = arrow::schema(fields); - - // 2. Open output file - auto result = fileSystem_->OpenOutputStream(schema_.paths[0]); - if (!result.ok()) { - return neug::Status( - neug::StatusCode::ERR_IO_ERROR, - "Failed to open output file: " + result.status().ToString()); - } - auto outfile = result.ValueOrDie(); - - // 3. Create Parquet writer with options - auto properties = buildWriterProperties(schema_.options); - - auto writer_result = ::parquet::arrow::FileWriter::Open( - *arrow_schema, arrow::default_memory_pool(), outfile, properties); - if (!writer_result.ok()) { - return neug::Status(neug::StatusCode::ERR_IO_ERROR, - "Failed to create Parquet writer: " + - writer_result.status().ToString()); - } - auto writer = std::move(writer_result.ValueOrDie()); - - // 4. Convert protobuf Arrays to Arrow Arrays std::vector> arrays; - for (int i = 0; i < num_columns; ++i) { - const auto& proto_array = table->arrays(i); - auto arrow_type = arrow_schema->field(i)->type(); - auto arrow_array = - protoArrayToArrowArray(proto_array, arrow_type, table->row_count()); - arrays.push_back(arrow_array); - } - - // 5. Create Arrow Table and write - auto arrow_table = arrow::Table::Make(arrow_schema, arrays); - - auto write_status = - writer->WriteTable(*arrow_table, arrow_table->num_rows()); - if (!write_status.ok()) { - return neug::Status( - neug::StatusCode::ERR_IO_ERROR, - "Failed to write Parquet table: " + write_status.ToString()); - } - - // 6. Close writer to flush and write footer - auto close_status = writer->Close(); - if (!close_status.ok()) { - return neug::Status( - neug::StatusCode::ERR_IO_ERROR, - "Failed to close Parquet writer: " + close_status.ToString()); - } - - // 7. Close output stream - auto outfile_close_status = outfile->Close(); - if (!outfile_close_status.ok()) { - return neug::Status( - neug::StatusCode::ERR_IO_ERROR, - "Failed to close output stream: " + outfile_close_status.ToString()); + arrays.reserve(fields.size()); + for (size_t i = 0; i < chunk.col_num(); ++i) { + if (chunk.columns[i] == nullptr) { + continue; + } + arrays.push_back(contextColumnToArrowArray( + chunk.columns[i], + arrow_schema->field(static_cast(arrays.size()))->type())); } - return neug::Status::OK(); + return writeArrowTable(arrow_schema, arrays, schema_, fileSystem_); } catch (const std::exception& e) { return neug::Status( neug::StatusCode::ERR_IO_ERROR, @@ -462,7 +535,6 @@ neug::Status ArrowParquetExportWriter::writeTable(const QueryResponse* table) { namespace function { -// Export function execution static execution::Context parquetExecFunc( neug::execution::Context& ctx, reader::FileSchema& schema, const std::shared_ptr& entry_schema, @@ -478,16 +550,19 @@ static execution::Context parquetExecFunc( auto writer = std::make_shared( schema, std::move(arrowFs), entry_schema); - auto status = writer->write(ctx, graph); - if (!status.ok()) { - THROW_IO_EXCEPTION("Parquet export failed: " + status.ToString()); + auto source_types = ctx.column_types(); + auto chunks = execution::Sink::materialize_for_export(ctx, graph); + for (const auto& chunk : chunks) { + auto status = writer->write(chunk, source_types); + if (!status.ok()) { + THROW_IO_EXCEPTION("Parquet export failed: " + status.ToString()); + } } LOG(INFO) << "[Parquet Export] Export completed successfully"; ctx.clear(); return ctx; } -// Bind function static std::unique_ptr bindFunc( ExportFuncBindInput& bindInput) { return std::make_unique( diff --git a/extension/parquet/tests/parquet_export_test_utils.h b/extension/parquet/tests/parquet_export_test_utils.h new file mode 100644 index 000000000..aeb79b111 --- /dev/null +++ b/extension/parquet/tests/parquet_export_test_utils.h @@ -0,0 +1,252 @@ +/** + * Copyright 2020 Alibaba Group Holding Limited. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include +#include + +#include "neug/common/columns/list_columns.h" +#include "neug/common/columns/struct_columns.h" +#include "neug/common/columns/value_columns.h" +#include "neug/common/types.h" +#include "neug/common/types/property_types.h" +#include "neug/common/types/value.h" +#include "neug/generated/proto/response/response.pb.h" + +namespace neug { +namespace test { + +inline bool protoValueValid(const std::string& validity, int row_idx) { + return validity.empty() || + ((static_cast(validity[static_cast(row_idx) >> 3]) >> + (row_idx & 7)) & + 1); +} + +inline std::shared_ptr protoArrayToColumn(const Array& arr) { + switch (arr.typed_array_case()) { + case Array::kBoolArray: { + const auto& typed = arr.bool_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kInt32Array: { + const auto& typed = arr.int32_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kInt64Array: { + const auto& typed = arr.int64_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kFloatArray: { + const auto& typed = arr.float_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kDoubleArray: { + const auto& typed = arr.double_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kStringArray: { + const auto& typed = arr.string_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kDateArray: { + const auto& typed = arr.date_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + Date date; + date.from_timestamp(typed.values(i)); + builder->push_back_opt(date); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kTimestampArray: { + const auto& typed = arr.timestamp_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + int64_t raw = typed.values(i); + int64_t millis = raw >= 100000000000000LL ? raw / 1000 : raw; + builder->push_back_opt(DateTime(millis)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kListArray: { + const auto& typed = arr.list_array(); + auto elements_col = protoArrayToColumn(typed.elements()); + const auto& child_type = elements_col->elem_type(); + ListColumnBuilder builder(child_type); + for (int row = 0; row + 1 < typed.offsets_size(); ++row) { + int32_t start = typed.offsets(row); + int32_t end = typed.offsets(row + 1); + std::vector values; + values.reserve(static_cast(end - start)); + for (int i = start; i < end; ++i) { + values.push_back(elements_col->get_elem(i)); + } + builder.push_back_elem(Value::LIST(child_type, std::move(values))); + } + return builder.finish(); + } + case Array::kStructArray: { + const auto& typed = arr.struct_array(); + std::vector child_types; + std::vector> field_cols; + field_cols.reserve(typed.fields_size()); + child_types.reserve(typed.fields_size()); + for (int i = 0; i < typed.fields_size(); ++i) { + field_cols.push_back(protoArrayToColumn(typed.fields(i))); + child_types.push_back(field_cols.back()->elem_type()); + } + auto struct_type = DataType::Struct(child_types); + StructColumnBuilder builder(struct_type); + const int num_rows = + field_cols.empty() ? 0 : static_cast(field_cols[0]->size()); + for (int row = 0; row < num_rows; ++row) { + if (!typed.validity().empty() && + !protoValueValid(typed.validity(), row)) { + builder.push_back_null(); + continue; + } + std::vector values; + values.reserve(field_cols.size()); + for (const auto& field_col : field_cols) { + values.push_back(field_col->get_elem(row)); + } + builder.push_back_elem(Value::STRUCT(struct_type, std::move(values))); + } + return builder.finish(); + } + case Array::kVertexArray: { + const auto& typed = arr.vertex_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kEdgeArray: { + const auto& typed = arr.edge_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + case Array::kPathArray: { + const auto& typed = arr.path_array(); + bool has_nulls = !typed.validity().empty(); + auto builder = std::make_shared>(has_nulls); + for (int i = 0; i < typed.values_size(); ++i) { + if (protoValueValid(typed.validity(), i)) { + builder->push_back_opt(typed.values(i)); + } else { + builder->push_back_null(); + } + } + return builder->finish(); + } + default: + return nullptr; + } +} + +inline DataChunk queryResponseToDataChunk(const QueryResponse& response) { + DataChunk chunk; + for (int i = 0; i < response.arrays_size(); ++i) { + auto column = protoArrayToColumn(response.arrays(i)); + if (column != nullptr) { + chunk.set(i, std::move(column)); + } + } + return chunk; +} + +} // namespace test +} // namespace neug diff --git a/extension/parquet/tests/parquet_test.cc b/extension/parquet/tests/parquet_test.cc index 334b4b695..a432a0b81 100644 --- a/extension/parquet/tests/parquet_test.cc +++ b/extension/parquet/tests/parquet_test.cc @@ -37,6 +37,7 @@ #include "../../extension/parquet/include/parquet_export_function.h" #include "../../extension/parquet/include/parquet_options.h" #include "neug/generated/proto/response/response.pb.h" +#include "parquet_export_test_utils.h" namespace neug { namespace test { @@ -454,7 +455,7 @@ TEST_F(ParquetTest, TestTypeMapping_StringToLargeUtf8) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Verify string column type auto col1 = ctx.chunk(0).columns()[1]; @@ -506,7 +507,7 @@ TEST_F(ParquetTest, TestTypeMapping_PreserveNumericTypes) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 4); EXPECT_EQ(ctx.row_num(), 1); @@ -578,7 +579,7 @@ TEST_F(ParquetTest, TestIntegration_ColumnPruning) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Verify extension translates projectColumns to Arrow projection // Should have 3 columns (id, score, grade - "name" is excluded) @@ -649,7 +650,7 @@ TEST_F(ParquetTest, TestIntegration_FilterPushdown) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Verify extension translates Neug filter to Arrow filter EXPECT_EQ(ctx.col_num(), 2); @@ -682,7 +683,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadMode) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_GT(ctx.chunk_num(), 0); // batch mode: data materialized into chunks EXPECT_GT(ctx.col_num(), 0) << "Extension should materialize data into " @@ -697,7 +698,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadMode) { auto reader2 = createParquetReader(sharedState2); auto localState2 = std::make_shared(); execution::Context ctx2; - reader2->read(localState2, ctx2); + ctx2.append_chunks(reader2->read(localState2)); auto col0_2 = ctx2.chunk(0).columns()[0]; EXPECT_EQ(col0_2->column_type(), ContextColumnType::kValue) @@ -762,7 +763,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilter) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Arrow scanner applies filter in both batch and full modes EXPECT_EQ(ctx.col_num(), 2); @@ -846,7 +847,7 @@ TEST_F(ParquetTest, TestIntegration_BatchReadWithFilterAndProjection) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Arrow scanner applies both filter and projection in batch mode EXPECT_EQ(ctx.col_num(), 2) @@ -932,7 +933,7 @@ TEST_F(ParquetTest, TestIntegration_CombinedFilterAndProjection) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Verify extension correctly combines filter and projection EXPECT_EQ(ctx.col_num(), 3) @@ -992,7 +993,7 @@ TEST_F(ParquetTest, TestMultiFile_ExplicitPaths) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 1); EXPECT_EQ(ctx.row_num(), 30) << "Extension should correctly read and " @@ -1055,7 +1056,7 @@ TEST_F(ParquetTest, TestParquetExportWriter) { entry_schema); // Write the response - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet: " << status.ToString(); // Verify file was created @@ -1070,7 +1071,7 @@ TEST_F(ParquetTest, TestParquetExportWriter) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 3); @@ -1116,7 +1117,7 @@ TEST_F(ParquetTest, TestParquetExportWithNulls) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with nulls: " << status.ToString(); @@ -1130,7 +1131,7 @@ TEST_F(ParquetTest, TestParquetExportWithNulls) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 2); EXPECT_EQ(ctx.row_num(), 3); @@ -1202,7 +1203,7 @@ TEST_F(ParquetTest, TestParquetExportMultipleTypes) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with multiple types: " << status.ToString(); @@ -1247,7 +1248,7 @@ TEST_F(ParquetTest, TestParquetExportLargeDataset) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write large Parquet: " << status.ToString(); @@ -1300,7 +1301,7 @@ TEST_F(ParquetTest, TestParquetExportWithCompressionOptions) { writer::ArrowParquetExportWriter writer_zstd(file_schema_zstd, file_system, entry_schema); - auto status = writer_zstd.writeTable(&response); + auto status = writer_zstd.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write ZSTD Parquet: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path_zstd)); @@ -1316,7 +1317,7 @@ TEST_F(ParquetTest, TestParquetExportWithCompressionOptions) { writer::ArrowParquetExportWriter writer_none(file_schema_none, file_system, entry_schema); - status = writer_none.writeTable(&response); + status = writer_none.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write uncompressed Parquet: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path_none)); @@ -1335,7 +1336,7 @@ TEST_F(ParquetTest, TestParquetExportWithCompressionOptions) { auto reader_zstd = createParquetReader(sharedState_zstd); auto localState_zstd = std::make_shared(); execution::Context ctx_zstd; - reader_zstd->read(localState_zstd, ctx_zstd); + ctx_zstd.append_chunks(reader_zstd->read(localState_zstd)); EXPECT_EQ(ctx_zstd.row_num(), 100); auto sharedState_none = createSharedState( @@ -1345,7 +1346,7 @@ TEST_F(ParquetTest, TestParquetExportWithCompressionOptions) { auto reader_none = createParquetReader(sharedState_none); auto localState_none = std::make_shared(); execution::Context ctx_none; - reader_none->read(localState_none, ctx_none); + ctx_none.append_chunks(reader_none->read(localState_none)); EXPECT_EQ(ctx_none.row_num(), 100); } @@ -1379,7 +1380,7 @@ TEST_F(ParquetTest, TestParquetExportWithUnsupportedCompression) { entry_schema); // Should fail due to unsupported codec - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); EXPECT_FALSE(status.ok()) << "Expected failure for unsupported codec, but got OK"; } @@ -1415,7 +1416,7 @@ TEST_F(ParquetTest, TestParquetExportWithRowGroupSize) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with row_group_size: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1428,7 +1429,7 @@ TEST_F(ParquetTest, TestParquetExportWithRowGroupSize) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.row_num(), 100); } @@ -1485,7 +1486,7 @@ TEST_F(ParquetTest, TestParquetExportWithDictionaryEncoding) { writer::ArrowParquetExportWriter writer_dict(file_schema_dict, file_system, entry_schema); - auto status = writer_dict.writeTable(&response); + auto status = writer_dict.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with dictionary encoding: " << status.ToString(); @@ -1503,7 +1504,7 @@ TEST_F(ParquetTest, TestParquetExportWithDictionaryEncoding) { writer::ArrowParquetExportWriter writer_nodict(file_schema_nodict, file_system, entry_schema); - status = writer_nodict.writeTable(&response); + status = writer_nodict.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet without dictionary encoding: " << status.ToString(); @@ -1526,7 +1527,7 @@ TEST_F(ParquetTest, TestParquetExportWithDictionaryEncoding) { auto reader_dict = createParquetReader(sharedState_dict); auto localState_dict = std::make_shared(); execution::Context ctx_dict; - reader_dict->read(localState_dict, ctx_dict); + ctx_dict.append_chunks(reader_dict->read(localState_dict)); EXPECT_EQ(ctx_dict.row_num(), num_rows); auto sharedState_nodict = createSharedState( @@ -1536,7 +1537,7 @@ TEST_F(ParquetTest, TestParquetExportWithDictionaryEncoding) { auto reader_nodict = createParquetReader(sharedState_nodict); auto localState_nodict = std::make_shared(); execution::Context ctx_nodict; - reader_nodict->read(localState_nodict, ctx_nodict); + ctx_nodict.append_chunks(reader_nodict->read(localState_nodict)); EXPECT_EQ(ctx_nodict.row_num(), num_rows); } @@ -1594,7 +1595,7 @@ TEST_F(ParquetTest, TestParquetExportWithDateAndTimestamp) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with date/timestamp: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1608,7 +1609,7 @@ TEST_F(ParquetTest, TestParquetExportWithDateAndTimestamp) { auto reader = createParquetReader(sharedState); auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.row_num(), num_rows); } @@ -1670,7 +1671,7 @@ TEST_F(ParquetTest, TestParquetExportWithListType) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with list type: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1746,7 +1747,7 @@ TEST_F(ParquetTest, TestParquetExportWithListOfStrings) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with list: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1814,7 +1815,7 @@ TEST_F(ParquetTest, TestParquetExportWithStructType) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with struct type: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1869,7 +1870,7 @@ TEST_F(ParquetTest, TestParquetExportWithVertexType) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with vertex type: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1924,7 +1925,7 @@ TEST_F(ParquetTest, TestParquetExportWithEdgeType) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with edge type: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -1976,7 +1977,7 @@ TEST_F(ParquetTest, TestParquetExportWithPathType) { writer::ArrowParquetExportWriter writer(file_schema, file_system, entry_schema); - auto status = writer.writeTable(&response); + auto status = writer.write(queryResponseToDataChunk(response)); ASSERT_TRUE(status.ok()) << "Failed to write Parquet with path type: " << status.ToString(); ASSERT_TRUE(std::filesystem::exists(export_path)); @@ -2005,7 +2006,7 @@ TEST_F(ParquetTest, TestParquetNonExistentColumnThrows) { auto localState = std::make_shared(); execution::Context ctx; - EXPECT_THROW(reader->read(localState, ctx), + EXPECT_THROW(ctx.append_chunks(reader->read(localState)), exception::SchemaMismatchException); } diff --git a/extension/pattern_matching/include/fastest_lib/src/DataStructure/graph.h b/extension/pattern_matching/include/fastest_lib/src/DataStructure/graph.h index 0a3dc6912..ce89097a9 100644 --- a/extension/pattern_matching/include/fastest_lib/src/DataStructure/graph.h +++ b/extension/pattern_matching/include/fastest_lib/src/DataStructure/graph.h @@ -35,7 +35,7 @@ // Include neug types #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" // Forward declarations removed - gbi types not used diff --git a/include/neug/common/columns/edge_columns.h b/include/neug/common/columns/edge_columns.h index cf7428e9a..948258d85 100644 --- a/include/neug/common/columns/edge_columns.h +++ b/include/neug/common/columns/edge_columns.h @@ -17,8 +17,7 @@ #include "neug/common/columns/columns_utils.h" #include "neug/common/types/graph_types.h" #include "neug/common/types/i_context_column.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/common/columns/value_columns.h b/include/neug/common/columns/value_columns.h index 033975b11..19b95d6fc 100644 --- a/include/neug/common/columns/value_columns.h +++ b/include/neug/common/columns/value_columns.h @@ -15,7 +15,7 @@ #pragma once #include "neug/common/columns/columns_utils.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" #include "neug/utils/top_n_generator.h" namespace neug { diff --git a/include/neug/utils/property/default_value.h b/include/neug/common/types/default_value.h similarity index 100% rename from include/neug/utils/property/default_value.h rename to include/neug/common/types/default_value.h diff --git a/include/neug/common/types/graph_types.h b/include/neug/common/types/graph_types.h index 8d69731d8..c9ed75393 100644 --- a/include/neug/common/types/graph_types.h +++ b/include/neug/common/types/graph_types.h @@ -23,7 +23,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/common/types/i_context_column.h b/include/neug/common/types/i_context_column.h index 12bf09bec..a0453bbea 100644 --- a/include/neug/common/types/i_context_column.h +++ b/include/neug/common/types/i_context_column.h @@ -23,7 +23,7 @@ #include "neug/common/types/value.h" #include "glog/logging.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/utils/property/types.h b/include/neug/common/types/property_types.h similarity index 100% rename from include/neug/utils/property/types.h rename to include/neug/common/types/property_types.h diff --git a/include/neug/compiler/function/export/json_export_function.h b/include/neug/compiler/function/export/json_export_function.h index 3f2495f48..0775364de 100644 --- a/include/neug/compiler/function/export/json_export_function.h +++ b/include/neug/compiler/function/export/json_export_function.h @@ -16,8 +16,6 @@ #pragma once -#include - #include "neug/compiler/function/export/export_function.h" #include "neug/utils/io/write/writer.h" #include "neug/utils/result.h" @@ -27,37 +25,6 @@ namespace neug { namespace writer { static constexpr const char* DEFAULT_JSON_NEWLINE = "\n"; -class JsonArrayStringFormatBuffer : public StringFormatBuffer { - public: - JsonArrayStringFormatBuffer(const neug::QueryResponse* response, - const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema); - ~JsonArrayStringFormatBuffer() = default; - void addValue(int rowIdx, int colIdx) override; - neug::Status flush(io::OutputStream& stream) override; - - private: - const reader::EntrySchema& entry_schema_; - rapidjson::Value current_line_; - rapidjson::Value buffer_; - rapidjson::Document document_; -}; - -class JsonLStringFormatBuffer : public StringFormatBuffer { - public: - JsonLStringFormatBuffer(const neug::QueryResponse* response, - const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema); - ~JsonLStringFormatBuffer() = default; - void addValue(int rowIdx, int colIdx) override; - neug::Status flush(io::OutputStream& stream) override; - - private: - const reader::EntrySchema& entry_schema_; - rapidjson::Value current_line_; - std::vector buffer_; - rapidjson::Document document_; -}; class JsonArrayExportWriter : public QueryExportWriter { public: @@ -67,7 +34,8 @@ class JsonArrayExportWriter : public QueryExportWriter { : QueryExportWriter(schema, std::move(entry_schema)) {} ~JsonArrayExportWriter() override = default; - neug::Status writeTable(const QueryResponse* table) override; + neug::Status write(const DataChunk& chunk, + const std::vector& source_types = {}) override; }; class JsonLExportWriter : public QueryExportWriter { @@ -78,7 +46,8 @@ class JsonLExportWriter : public QueryExportWriter { : QueryExportWriter(schema, std::move(entry_schema)) {} ~JsonLExportWriter() override = default; - neug::Status writeTable(const QueryResponse* table) override; + neug::Status write(const DataChunk& chunk, + const std::vector& source_types = {}) override; }; } // namespace writer diff --git a/include/neug/compiler/function/import/csv_read_function.h b/include/neug/compiler/function/import/csv_read_function.h index 9a55756f2..672644c24 100644 --- a/include/neug/compiler/function/import/csv_read_function.h +++ b/include/neug/compiler/function/import/csv_read_function.h @@ -128,7 +128,7 @@ struct CSVReadFunction { std::make_unique(state, std::move(optionsBuilder)); execution::Context ctx; auto localState = std::make_shared(); - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); return ctx; } diff --git a/include/neug/compiler/function/import/json_read_function.h b/include/neug/compiler/function/import/json_read_function.h index 5b2e8793c..1810d5726 100644 --- a/include/neug/compiler/function/import/json_read_function.h +++ b/include/neug/compiler/function/import/json_read_function.h @@ -61,7 +61,7 @@ struct JsonReadFunction { std::make_unique(state, std::move(optionsBuilder)); execution::Context ctx; auto localState = std::make_shared(); - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); return ctx; } @@ -127,7 +127,7 @@ struct JsonLReadFunction { std::make_unique(state, std::move(optionsBuilder)); execution::Context ctx; auto localState = std::make_shared(); - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); return ctx; } diff --git a/include/neug/execution/common/context.h b/include/neug/execution/common/context.h index 3d3e0047c..167eb739d 100644 --- a/include/neug/execution/common/context.h +++ b/include/neug/execution/common/context.h @@ -74,6 +74,11 @@ class Context { /// Appends a fully-formed ContextChunk to this Context. void append_chunk(ContextChunk&& chunk); + /// Appends multiple chunks (each with no head) to this Context. Convenience + /// for consuming the output of the IO read layer, which produces plain + /// DataChunks with no knowledge of Context. + void append_chunks(std::vector&& chunks); + /// Applies a chunk-level operation to every chunk in this Context. template neug::result apply_chunks(F&& func) { @@ -106,6 +111,10 @@ class Context { /// Returns the total number of rows across all chunks. size_t row_num() const; + /// Column logical types in tag_ids order (from the first chunk that has each + /// tag). + std::vector column_types() const; + std::vector tag_ids; private: diff --git a/include/neug/execution/common/operators/retrieve/sink.h b/include/neug/execution/common/operators/retrieve/sink.h index eabfa59ff..211701d64 100644 --- a/include/neug/execution/common/operators/retrieve/sink.h +++ b/include/neug/execution/common/operators/retrieve/sink.h @@ -14,6 +14,8 @@ */ #pragma once +#include "neug/common/types.h" +#include "neug/common/types/data_chunk.h" #include "neug/generated/proto/response/response.pb.h" #include "neug/main/query_result.h" namespace neug { @@ -27,6 +29,11 @@ class Context; class Sink { public: + /// Materialize graph-type columns to VARCHAR strings per Context chunk. + /// Returns one DataChunk per Context chunk without cross-chunk column fusion. + static std::vector materialize_for_export( + const Context& ctx, const StorageReadInterface& graph); + static void sink_results(const Context& ctx, const StorageReadInterface& graph, neug::QueryResponse* response); diff --git a/include/neug/execution/execute/ops/batch/batch_update_utils.h b/include/neug/execution/execute/ops/batch/batch_update_utils.h index 70ba91dec..90193ccf5 100644 --- a/include/neug/execution/execute/ops/batch/batch_update_utils.h +++ b/include/neug/execution/execute/ops/batch/batch_update_utils.h @@ -17,8 +17,8 @@ #include #include "neug/common/types/graph_types.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" -#include "neug/utils/property/types.h" namespace physical { class PropertyMapping; diff --git a/include/neug/execution/execute/ops/retrieve/scan_utils.h b/include/neug/execution/execute/ops/retrieve/scan_utils.h index a87994685..96be713aa 100644 --- a/include/neug/execution/execute/ops/retrieve/scan_utils.h +++ b/include/neug/execution/execute/ops/retrieve/scan_utils.h @@ -19,9 +19,9 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/execution/execute/operator.h" -#include "neug/utils/property/types.h" namespace algebra { class IndexPredicate; diff --git a/include/neug/execution/expression/special_predicates.h b/include/neug/execution/expression/special_predicates.h index bcc6112db..4c92e509d 100644 --- a/include/neug/execution/expression/special_predicates.h +++ b/include/neug/execution/expression/special_predicates.h @@ -17,10 +17,10 @@ #include "neug/execution/utils/pb_parse_utils.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/execution/common/params_map.h" #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" namespace common { class Expression; diff --git a/include/neug/main/neug_db.h b/include/neug/main/neug_db.h index 355a1a0a2..04261ba4d 100644 --- a/include/neug/main/neug_db.h +++ b/include/neug/main/neug_db.h @@ -22,6 +22,7 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/config.h" #include "neug/execution/execute/query_cache.h" #include "neug/generated/proto/plan/cypher_ddl.pb.h" @@ -36,7 +37,6 @@ #include "neug/transaction/insert_transaction.h" #include "neug/transaction/read_transaction.h" #include "neug/transaction/update_transaction.h" -#include "neug/utils/property/types.h" #include "neug/version.h" #define STRINGIFY(x) #x diff --git a/include/neug/storages/README.md b/include/neug/storages/README.md index 7de825533..fe20144b9 100644 --- a/include/neug/storages/README.md +++ b/include/neug/storages/README.md @@ -26,11 +26,11 @@ In `PropertyGraph`, schema is defined [here](./graph/schema.h). To initialize an In `PropertyGraph`, each vertex is assigned an internal ID starting from zero, which is a consecutive integer. This ID can uniquely identify a vertex and can also be used as an offset to quickly access data related to the vertex in property tables and adjacency lists. -The mapping between the external and internal IDs is maintained in [LFIndexer](../utils/id_indexer.h), which is implemented based on ska::flat_hash_map and supports concurrent, lock-free insertion of elements. +The mapping between the external and internal IDs is maintained in [LFIndexer](id_indexer.h), which is implemented based on ska::flat_hash_map and supports concurrent, lock-free insertion of elements. ### 3.2 Vertex properties storage -The properties of vertices are maintained in [Table](../utils/property/table.h) in a column-oriented manner. Storing data in a column-oriented manner makes it more compact, and users can access only the necessary columns as per their requirements, without introducing redundant memory/disk access. +The properties of vertices are maintained in [Table](../property/table.h) in a column-oriented manner. Storing data in a column-oriented manner makes it more compact, and users can access only the necessary columns as per their requirements, without introducing redundant memory/disk access. All the columns are implemented with `mmap`-ed memory. The `max_vertex_num` field of Schema will be used to ensure virtual memory region is large enough. diff --git a/include/neug/storages/allocators.h b/include/neug/storages/allocators.h index ea00dc2b3..e533eaf81 100644 --- a/include/neug/storages/allocators.h +++ b/include/neug/storages/allocators.h @@ -23,12 +23,12 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/config.h" #include "neug/storages/container/container_utils.h" #include "neug/storages/container/file_header.h" #include "neug/storages/container/mmap_container.h" #include "neug/utils/io/file/file_utils.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/storages/container/i_container.h b/include/neug/storages/container/i_container.h index 985889271..79ba773c9 100644 --- a/include/neug/storages/container/i_container.h +++ b/include/neug/storages/container/i_container.h @@ -18,8 +18,8 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/config.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/storages/csr/csr_base.h b/include/neug/storages/csr/csr_base.h index dbaed6c0a..29c6095cb 100644 --- a/include/neug/storages/csr/csr_base.h +++ b/include/neug/storages/csr/csr_base.h @@ -18,12 +18,12 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/allocators.h" #include "neug/storages/csr/csr_view.h" #include "neug/storages/csr/nbr.h" #include "neug/storages/module/module.h" -#include "neug/utils/property/types.h" #include diff --git a/include/neug/storages/csr/csr_view.h b/include/neug/storages/csr/csr_view.h index 659bfdb52..1be0dd49b 100644 --- a/include/neug/storages/csr/csr_view.h +++ b/include/neug/storages/csr/csr_view.h @@ -16,11 +16,11 @@ #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/csr/nbr.h" #include "neug/storages/csr/prefetch_utils.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/types.h" +#include "neug/storages/property/column.h" namespace neug { diff --git a/include/neug/storages/csr/immutable_csr.h b/include/neug/storages/csr/immutable_csr.h index bf9e51670..8fda8b859 100644 --- a/include/neug/storages/csr/immutable_csr.h +++ b/include/neug/storages/csr/immutable_csr.h @@ -21,12 +21,12 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/container/i_container.h" #include "neug/storages/csr/csr_base.h" #include "neug/storages/csr/csr_view.h" #include "neug/storages/csr/nbr.h" #include "neug/storages/module/type_name.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/storages/csr/mutable_csr.h b/include/neug/storages/csr/mutable_csr.h index 61a685502..610c78b6a 100644 --- a/include/neug/storages/csr/mutable_csr.h +++ b/include/neug/storages/csr/mutable_csr.h @@ -31,6 +31,7 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/allocators.h" #include "neug/storages/container/i_container.h" #include "neug/storages/csr/csr_base.h" @@ -38,7 +39,6 @@ #include "neug/storages/csr/nbr.h" #include "neug/storages/module/type_name.h" #include "neug/utils/io/file/file_utils.h" -#include "neug/utils/property/types.h" #include "neug/utils/spinlock.h" namespace neug { diff --git a/include/neug/storages/csr/nbr.h b/include/neug/storages/csr/nbr.h index 98716da57..b1686bdc8 100644 --- a/include/neug/storages/csr/nbr.h +++ b/include/neug/storages/csr/nbr.h @@ -14,9 +14,9 @@ */ #pragma once -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" +#include "neug/storages/property/column.h" +#include "neug/storages/property/table.h" namespace neug { diff --git a/include/neug/storages/graph/edge_table.h b/include/neug/storages/graph/edge_table.h index a32dd6007..d5785132c 100644 --- a/include/neug/storages/graph/edge_table.h +++ b/include/neug/storages/graph/edge_table.h @@ -23,16 +23,16 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/allocators.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/csr/csr_base.h" #include "neug/storages/csr/csr_view.h" #include "neug/storages/graph/schema.h" +#include "neug/storages/indexers.h" #include "neug/storages/module/module.h" -#include "neug/utils/indexers.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" +#include "neug/storages/property/table.h" namespace neug { diff --git a/include/neug/storages/graph/graph_interface.h b/include/neug/storages/graph/graph_interface.h index 3d0daf7b7..52ede0caf 100644 --- a/include/neug/storages/graph/graph_interface.h +++ b/include/neug/storages/graph/graph_interface.h @@ -18,11 +18,11 @@ #include #include "neug/common/types/container_types.h" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/graph/graph_view.h" #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/storages/graph/graph_view.h b/include/neug/storages/graph/graph_view.h index 0e1e49745..537802a6c 100644 --- a/include/neug/storages/graph/graph_view.h +++ b/include/neug/storages/graph/graph_view.h @@ -26,7 +26,7 @@ #include "neug/storages/graph/edge_table.h" #include "neug/storages/graph/schema.h" #include "neug/storages/graph/vertex_table.h" -#include "neug/utils/property/column.h" +#include "neug/storages/property/column.h" #include "neug/utils/result.h" namespace neug { diff --git a/include/neug/storages/graph/operation_params.h b/include/neug/storages/graph/operation_params.h index 6abd423b9..b77579dc7 100644 --- a/include/neug/storages/graph/operation_params.h +++ b/include/neug/storages/graph/operation_params.h @@ -18,8 +18,8 @@ #include #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" -#include "neug/utils/property/types.h" namespace neug { class InArchive; diff --git a/include/neug/storages/graph/property_graph.h b/include/neug/storages/graph/property_graph.h index 72cf88fac..b64dc8554 100644 --- a/include/neug/storages/graph/property_graph.h +++ b/include/neug/storages/graph/property_graph.h @@ -27,6 +27,7 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/allocators.h" #include "neug/storages/checkpoint.h" @@ -38,7 +39,6 @@ #include "neug/storages/graph/schema.h" #include "neug/storages/graph/vertex_table.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { diff --git a/include/neug/storages/graph/property_graph_cow_state.h b/include/neug/storages/graph/property_graph_cow_state.h index 4b448abe0..b96629a77 100644 --- a/include/neug/storages/graph/property_graph_cow_state.h +++ b/include/neug/storages/graph/property_graph_cow_state.h @@ -20,7 +20,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/storages/graph/schema.h b/include/neug/storages/graph/schema.h index d5e1cf092..ea08d8bc5 100644 --- a/include/neug/storages/graph/schema.h +++ b/include/neug/storages/graph/schema.h @@ -26,11 +26,11 @@ #include #include #include +#include "neug/common/types/default_value.h" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/utils/bitset.h" -#include "neug/utils/property/default_value.h" #include "neug/utils/property/property_definition.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" diff --git a/include/neug/storages/graph/vertex_table.h b/include/neug/storages/graph/vertex_table.h index 7a9ba701b..3fc6f649d 100644 --- a/include/neug/storages/graph/vertex_table.h +++ b/include/neug/storages/graph/vertex_table.h @@ -18,10 +18,10 @@ #include "neug/common/types/value.h" #include "neug/storages/graph/schema.h" #include "neug/storages/graph/vertex_timestamp.h" +#include "neug/storages/indexers.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module.h" -#include "neug/utils/indexers.h" -#include "neug/utils/property/table.h" +#include "neug/storages/property/table.h" namespace neug { diff --git a/include/neug/storages/graph/vertex_timestamp.h b/include/neug/storages/graph/vertex_timestamp.h index 6287cd8d3..ff6751037 100644 --- a/include/neug/storages/graph/vertex_timestamp.h +++ b/include/neug/storages/graph/vertex_timestamp.h @@ -20,9 +20,9 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/module/module.h" #include "neug/utils/likely.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/utils/id_indexer.h b/include/neug/storages/id_indexer.h similarity index 99% rename from include/neug/utils/id_indexer.h rename to include/neug/storages/id_indexer.h index 375ad26c9..fe0a4913f 100644 --- a/include/neug/utils/id_indexer.h +++ b/include/neug/storages/id_indexer.h @@ -34,15 +34,15 @@ limitations under the License. #include "glog/logging.h" #include "neug/common/columns/value_columns.h" #include "neug/common/types/i_context_column.h" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/container/container_utils.h" #include "neug/storages/container/i_container.h" #include "neug/storages/module/module.h" +#include "neug/storages/property/column.h" #include "neug/utils/bitset.h" #include "neug/utils/likely.h" #include "neug/utils/pb_utils.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" #include "neug/utils/string_view_vector.h" diff --git a/include/neug/utils/indexers.h b/include/neug/storages/indexers.h similarity index 89% rename from include/neug/utils/indexers.h rename to include/neug/storages/indexers.h index 4c3b453b5..9bec51326 100644 --- a/include/neug/utils/indexers.h +++ b/include/neug/storages/indexers.h @@ -14,9 +14,9 @@ limitations under the License. */ #pragma once -#include "neug/utils/id_indexer.h" +#include "neug/storages/id_indexer.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/storages/loader/abstract_property_graph_loader.h b/include/neug/storages/loader/abstract_property_graph_loader.h index 8a974bf67..ec469cea3 100644 --- a/include/neug/storages/loader/abstract_property_graph_loader.h +++ b/include/neug/storages/loader/abstract_property_graph_loader.h @@ -16,12 +16,12 @@ #include +#include "neug/common/types/property_types.h" #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" #include "neug/storages/loader/i_fragment_loader.h" #include "neug/storages/loader/loading_config.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { diff --git a/include/neug/storages/loader/i_fragment_loader.h b/include/neug/storages/loader/i_fragment_loader.h index ace752c4f..9920b268a 100644 --- a/include/neug/storages/loader/i_fragment_loader.h +++ b/include/neug/storages/loader/i_fragment_loader.h @@ -14,10 +14,10 @@ */ #pragma once +#include "neug/common/types/property_types.h" #include "neug/storages/csr/mutable_csr.h" #include "neug/storages/graph/schema.h" -#include "neug/utils/indexers.h" -#include "neug/utils/property/types.h" +#include "neug/storages/indexers.h" #include "neug/utils/result.h" namespace neug { diff --git a/include/neug/storages/loader/loading_config.h b/include/neug/storages/loader/loading_config.h index df2f47f31..815815829 100644 --- a/include/neug/storages/loader/loading_config.h +++ b/include/neug/storages/loader/loading_config.h @@ -26,8 +26,8 @@ #include #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/graph/schema.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace YAML { diff --git a/include/neug/storages/module/type_name.h b/include/neug/storages/module/type_name.h index 2fa00fbdd..e86350aaa 100644 --- a/include/neug/storages/module/type_name.h +++ b/include/neug/storages/module/type_name.h @@ -18,8 +18,8 @@ #include #include "neug/common/types.h" +#include "neug/common/types/property_types.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" namespace neug { @@ -31,7 +31,7 @@ namespace neug { * NEUG_REGISTER_TEMPLATE_MODULE. Add a specialization here whenever a new * leaf storage type is introduced. Types listed below correspond to the * registrations in: - * - src/utils/property/column.cc (TypedColumn) + * - src/storages/property/column.cc (TypedColumn) * - src/storages/csr/mutable_csr.cc (MutableCsr, SingleMutableCsr, * EmptyCsr) * - src/storages/csr/immutable_csr.cc (ImmutableCsr, diff --git a/include/neug/utils/property/array_column.h b/include/neug/storages/property/array_column.h similarity index 98% rename from include/neug/utils/property/array_column.h rename to include/neug/storages/property/array_column.h index 4ecca9be3..7e37f33d8 100644 --- a/include/neug/utils/property/array_column.h +++ b/include/neug/storages/property/array_column.h @@ -20,7 +20,7 @@ #include "neug/common/extra_type_info.h" #include "neug/common/types.h" #include "neug/common/types/value.h" -#include "neug/utils/property/column.h" +#include "neug/storages/property/column.h" namespace neug { diff --git a/include/neug/utils/property/column.h b/include/neug/storages/property/column.h similarity index 99% rename from include/neug/utils/property/column.h rename to include/neug/storages/property/column.h index f6c7f802c..e80157f3c 100644 --- a/include/neug/utils/property/column.h +++ b/include/neug/storages/property/column.h @@ -31,6 +31,7 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/config.h" #include "neug/storages/checkpoint.h" @@ -43,7 +44,6 @@ #include "neug/utils/exception/exception.h" #include "neug/utils/io/file/file_utils.h" #include "neug/utils/likely.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/out_archive.h" #include diff --git a/include/neug/utils/property/table.h b/include/neug/storages/property/table.h similarity index 97% rename from include/neug/utils/property/table.h rename to include/neug/storages/property/table.h index fa8bcd086..1ce71081c 100644 --- a/include/neug/utils/property/table.h +++ b/include/neug/storages/property/table.h @@ -21,13 +21,13 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/config.h" #include "neug/storages/checkpoint.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/module/module.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/types.h" +#include "neug/storages/property/column.h" namespace neug { diff --git a/include/neug/transaction/compact_transaction.h b/include/neug/transaction/compact_transaction.h index 5c011b881..6c72ae7d6 100644 --- a/include/neug/transaction/compact_transaction.h +++ b/include/neug/transaction/compact_transaction.h @@ -14,8 +14,8 @@ */ #pragma once +#include "neug/common/types/property_types.h" #include "neug/storages/graph_snapshot_store.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" namespace neug { diff --git a/include/neug/transaction/insert_transaction.h b/include/neug/transaction/insert_transaction.h index 8b271d364..586627d40 100644 --- a/include/neug/transaction/insert_transaction.h +++ b/include/neug/transaction/insert_transaction.h @@ -23,13 +23,13 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/allocators.h" #include "neug/storages/graph/graph_interface.h" #include "neug/storages/graph/graph_stats.h" #include "neug/storages/graph/graph_view.h" #include "neug/storages/graph_snapshot_store.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" namespace neug { diff --git a/include/neug/transaction/read_transaction.h b/include/neug/transaction/read_transaction.h index c0a5458b8..40662ae20 100644 --- a/include/neug/transaction/read_transaction.h +++ b/include/neug/transaction/read_transaction.h @@ -29,6 +29,7 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/csr/immutable_csr.h" #include "neug/storages/csr/mutable_csr.h" #include "neug/storages/csr/nbr.h" @@ -37,10 +38,9 @@ #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" #include "neug/storages/graph_snapshot_store.h" +#include "neug/storages/property/column.h" +#include "neug/storages/property/table.h" #include "neug/transaction/transaction_utils.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/transaction/transaction_utils.h b/include/neug/transaction/transaction_utils.h index f0d8f3c07..9fe328b91 100644 --- a/include/neug/transaction/transaction_utils.h +++ b/include/neug/transaction/transaction_utils.h @@ -15,9 +15,9 @@ #pragma once #include "glog/logging.h" +#include "neug/common/types/property_types.h" #include "neug/storages/graph/property_graph.h" #include "neug/utils/likely.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" diff --git a/include/neug/transaction/update_transaction.h b/include/neug/transaction/update_transaction.h index e35fb1eb4..4741e5710 100644 --- a/include/neug/transaction/update_transaction.h +++ b/include/neug/transaction/update_transaction.h @@ -24,6 +24,7 @@ #include #include "flat_hash_map/flat_hash_map.hpp" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/execution/execute/query_cache.h" #include "neug/storages/allocators.h" @@ -34,10 +35,9 @@ #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/property_graph_cow_state.h" #include "neug/storages/graph_snapshot_store.h" +#include "neug/storages/property/table.h" #include "neug/transaction/transaction_utils.h" #include "neug/transaction/wal/wal_builder.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/transaction/wal/wal.h b/include/neug/transaction/wal/wal.h index 74bd75397..0043215e1 100644 --- a/include/neug/transaction/wal/wal.h +++ b/include/neug/transaction/wal/wal.h @@ -21,10 +21,10 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/graph/operation_params.h" #include "neug/transaction/transaction_utils.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" diff --git a/include/neug/transaction/wal/wal_builder.h b/include/neug/transaction/wal/wal_builder.h index 73b76de6b..c5f57e864 100644 --- a/include/neug/transaction/wal/wal_builder.h +++ b/include/neug/transaction/wal/wal_builder.h @@ -18,11 +18,11 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/graph/operation_params.h" #include "neug/transaction/transaction_utils.h" #include "neug/transaction/wal/wal.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" namespace neug { diff --git a/include/neug/utils/io/read/common/type_converter.h b/include/neug/utils/io/read/common/type_converter.h index 5f0877c36..236c72cdc 100644 --- a/include/neug/utils/io/read/common/type_converter.h +++ b/include/neug/utils/io/read/common/type_converter.h @@ -17,8 +17,8 @@ #include +#include "neug/common/types/property_types.h" #include "neug/generated/proto/plan/basic_type.pb.h" -#include "neug/utils/property/types.h" namespace neug { namespace reader { diff --git a/include/neug/utils/io/read/csv/csv_read_config.h b/include/neug/utils/io/read/csv/csv_read_config.h index cf8aba834..02f910324 100644 --- a/include/neug/utils/io/read/csv/csv_read_config.h +++ b/include/neug/utils/io/read/csv/csv_read_config.h @@ -19,7 +19,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/utils/io/read/csv/csv_reader.h b/include/neug/utils/io/read/csv/csv_reader.h index bf4f66d74..d8d161951 100644 --- a/include/neug/utils/io/read/csv/csv_reader.h +++ b/include/neug/utils/io/read/csv/csv_reader.h @@ -18,7 +18,7 @@ #include #include -#include "neug/execution/common/context.h" +#include "neug/common/types/data_chunk.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 +28,6 @@ namespace neug { class IDataChunkSupplier; -namespace execution { -class Context; -} - namespace reader { class CsvReader { @@ -40,18 +36,16 @@ class CsvReader { std::unique_ptr optionsBuilder); ~CsvReader(); - void read(std::shared_ptr localState, - execution::Context& ctx); + std::vector read(std::shared_ptr localState); result> inferSchema(); private: - void full_read( - const std::vector>& suppliers, - execution::Context& output, const CsvReadConfig& output_config); - void batch_read( + std::vector full_read( const std::vector>& suppliers, - execution::Context& output); + const CsvReadConfig& output_config); + std::vector batch_read( + const std::vector>& suppliers); std::shared_ptr sharedState_; std::unique_ptr optionsBuilder_; diff --git a/include/neug/utils/io/read/json/json_read_config.h b/include/neug/utils/io/read/json/json_read_config.h index f677564c2..573496077 100644 --- a/include/neug/utils/io/read/json/json_read_config.h +++ b/include/neug/utils/io/read/json/json_read_config.h @@ -19,7 +19,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/utils/io/read/json/json_reader.h b/include/neug/utils/io/read/json/json_reader.h index b88011352..ace75d775 100644 --- a/include/neug/utils/io/read/json/json_reader.h +++ b/include/neug/utils/io/read/json/json_reader.h @@ -18,7 +18,7 @@ #include #include -#include "neug/execution/common/context.h" +#include "neug/common/types/data_chunk.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 +28,6 @@ namespace neug { class IDataChunkSupplier; -namespace execution { -class Context; -} - namespace reader { class JsonReader { @@ -40,18 +36,16 @@ class JsonReader { std::unique_ptr optionsBuilder); ~JsonReader(); - void read(std::shared_ptr localState, - execution::Context& ctx); + std::vector read(std::shared_ptr localState); result> inferSchema(); private: - void full_read( - const std::vector>& suppliers, - execution::Context& output, const JsonReadConfig& output_config); - void batch_read( + std::vector full_read( const std::vector>& suppliers, - execution::Context& output); + const JsonReadConfig& output_config); + std::vector batch_read( + const std::vector>& suppliers); std::shared_ptr sharedState_; std::unique_ptr optionsBuilder_; diff --git a/include/neug/utils/io/write/writer.h b/include/neug/utils/io/write/writer.h index 399bad15f..467dd160f 100644 --- a/include/neug/utils/io/write/writer.h +++ b/include/neug/utils/io/write/writer.h @@ -18,10 +18,8 @@ #include #include -#include "neug/execution/common/context.h" -#include "neug/execution/execute/ops/batch/batch_update_utils.h" -#include "neug/generated/proto/response/response.pb.h" -#include "neug/storages/graph/graph_interface.h" +#include "neug/common/types.h" +#include "neug/common/types/data_chunk.h" #include "neug/utils/io/read/common/options.h" #include "neug/utils/io/read/common/schema.h" #include "neug/utils/io/stream/output_stream.h" @@ -54,71 +52,51 @@ class ExportWriter { virtual ~ExportWriter() = default; - virtual neug::Status write(const execution::Context& context, - const StorageReadInterface& graph) = 0; + virtual neug::Status write( + const DataChunk& chunk, + const std::vector& source_types = {}) = 0; protected: const reader::FileSchema& schema_; std::shared_ptr entry_schema_; }; -class StringFormatBuffer { - public: - StringFormatBuffer(const neug::QueryResponse* response, - const reader::FileSchema& schema) - : response_(response), schema_(schema) {} - ~StringFormatBuffer() {} - virtual void addValue(int rowIdx, int colIdx) = 0; - virtual neug::Status flush(io::OutputStream& stream) = 0; - static bool validateIndex(const neug::QueryResponse* response, int rowIdx, - int colIdx); - static bool validateProtoValue(const std::string& validity, int rowIdx); - - protected: - const neug::QueryResponse* response_; - const reader::FileSchema& schema_; -}; - struct BinaryData { std::unique_ptr data; uint64_t size = 0; }; -class CSVStringFormatBuffer : public StringFormatBuffer { +class DataChunkCSVStringFormatBuffer { public: - CSVStringFormatBuffer(const neug::QueryResponse* response, - const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema); - ~CSVStringFormatBuffer() {} - void addValue(int rowIdx, int colIdx) override; + DataChunkCSVStringFormatBuffer(const DataChunk& chunk, + const reader::FileSchema& schema, + const reader::EntrySchema& entry_schema); + ~DataChunkCSVStringFormatBuffer() = default; + + void addValue(size_t row_idx, size_t col_idx); void addHeader(); - neug::Status flush(io::OutputStream& stream) override; + neug::Status flush(io::OutputStream& stream); private: + neug::Status formatValueToStr(const Value& value, size_t row_idx); + void writeWithEscapes(char* toEscape, char escape, const std::string& str); + void write(const uint8_t* buffer, uint64_t len); + + const DataChunk& chunk_; + const reader::FileSchema& schema_; + const reader::EntrySchema& entry_schema_; BinaryData blob_; size_t capacity_; uint8_t* data_; - const reader::EntrySchema& entry_schema_; - // Cached WriteOptions resolved once at construction to avoid per-cell lookup. bool has_header_; char delimiter_; bool ignore_errors_; char escape_char_; char quote_char_; - private: - // write the current value to string buffer, return error status if value is - // invalid - neug::Status formatValueToStr(const neug::Array& arr, int rowIdx); - void writeWithEscapes(char* toEscape, char escape, const std::string& str); - void write(const uint8_t* buffer, uint64_t len); - - private: static constexpr const char* DEFAULT_CSV_NEWLINE = "\n"; static constexpr const char* DEFAULT_NULL_STR = ""; static constexpr size_t DEFAULT_CAPACITY = 64; - static constexpr const char* LIST_ARRAY_CHAR = "[]"; - static constexpr const char* COMMA_CHAR = ","; }; class QueryExportWriter : public ExportWriter { @@ -128,11 +106,6 @@ class QueryExportWriter : public ExportWriter { std::shared_ptr entry_schema = nullptr) : ExportWriter(schema, std::move(entry_schema)) {} ~QueryExportWriter() override = default; - - neug::Status write(const execution::Context& context, - const StorageReadInterface& graph) override; - - virtual neug::Status writeTable(const QueryResponse* table) = 0; }; class CsvQueryExportWriter : public QueryExportWriter { @@ -143,7 +116,8 @@ class CsvQueryExportWriter : public QueryExportWriter { : QueryExportWriter(schema, std::move(entry_schema)) {} ~CsvQueryExportWriter() override = default; - neug::Status writeTable(const QueryResponse* table) override; + neug::Status write(const DataChunk& chunk, + const std::vector& source_types = {}) override; }; } // namespace writer diff --git a/include/neug/utils/pb_utils.h b/include/neug/utils/pb_utils.h index 9a1da2c66..5aeb47593 100644 --- a/include/neug/utils/pb_utils.h +++ b/include/neug/utils/pb_utils.h @@ -19,12 +19,12 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/generated/proto/plan/basic_type.pb.h" #include "neug/generated/proto/plan/cypher_ddl.pb.h" #include "neug/generated/proto/plan/physical.pb.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" #include "neug/utils/service_utils.h" diff --git a/include/neug/utils/serialization/in_archive.h b/include/neug/utils/serialization/in_archive.h index 04722adaf..5b44f8d47 100644 --- a/include/neug/utils/serialization/in_archive.h +++ b/include/neug/utils/serialization/in_archive.h @@ -20,7 +20,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/utils/serialization/out_archive.h b/include/neug/utils/serialization/out_archive.h index 2281616aa..b1f01cc37 100644 --- a/include/neug/utils/serialization/out_archive.h +++ b/include/neug/utils/serialization/out_archive.h @@ -20,7 +20,7 @@ limitations under the License. #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/utils/service_utils.h b/include/neug/utils/service_utils.h index 3b5c3adf4..2f89e057f 100644 --- a/include/neug/utils/service_utils.h +++ b/include/neug/utils/service_utils.h @@ -53,8 +53,8 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/include/neug/utils/string_utils.h b/include/neug/utils/string_utils.h index eac535763..35867c059 100644 --- a/include/neug/utils/string_utils.h +++ b/include/neug/utils/string_utils.h @@ -27,7 +27,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/src/utils/property/default_value.cc b/src/common/types/default_value.cc similarity index 97% rename from src/utils/property/default_value.cc rename to src/common/types/default_value.cc index 2bdbe9a7e..448284644 100644 --- a/src/utils/property/default_value.cc +++ b/src/common/types/default_value.cc @@ -13,7 +13,7 @@ * limitations under the License. */ -#include "neug/utils/property/default_value.h" +#include "neug/common/types/default_value.h" #include "neug/common/extra_type_info.h" #include "neug/common/types/value.h" diff --git a/src/common/types/graph_types.cc b/src/common/types/graph_types.cc index d9b7d7e6e..9991969bf 100644 --- a/src/common/types/graph_types.cc +++ b/src/common/types/graph_types.cc @@ -16,7 +16,7 @@ #include "neug/common/types/graph_types.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { int64_t encode_unique_vertex_id(label_t label_id, vid_t vid) { diff --git a/src/utils/property/types.cc b/src/common/types/property_types.cc similarity index 99% rename from src/utils/property/types.cc rename to src/common/types/property_types.cc index 3d2528836..0fbfdea2d 100644 --- a/src/utils/property/types.cc +++ b/src/common/types/property_types.cc @@ -13,7 +13,7 @@ * limitations under the License. */ -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" #include #include @@ -30,8 +30,6 @@ #include #include "date/date.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" diff --git a/src/compiler/function/csv_export_function.cpp b/src/compiler/function/csv_export_function.cpp index 0e044e553..c9b25ba37 100644 --- a/src/compiler/function/csv_export_function.cpp +++ b/src/compiler/function/csv_export_function.cpp @@ -22,6 +22,7 @@ #include "neug/compiler/function/export/export_function.h" #include "neug/compiler/main/metadata_registry.h" +#include "neug/execution/common/operators/retrieve/sink.h" #include "neug/utils/io/write/writer.h" namespace neug { @@ -62,12 +63,16 @@ execution::Context writeExecFunc( convertFileSchemaOptions(schema); auto writer = std::make_shared( schema, entry_schema); - auto status = writer->write(ctx, graph); - if (!status.ok()) { - if (status.error_code() == StatusCode::ERR_PERMISSION) { - THROW_PERMISSION_DENIED("Export failed: " + status.ToString()); + auto source_types = ctx.column_types(); + auto chunks = neug::execution::Sink::materialize_for_export(ctx, graph); + for (const auto& chunk : chunks) { + auto status = writer->write(chunk, source_types); + if (!status.ok()) { + if (status.error_code() == StatusCode::ERR_PERMISSION) { + THROW_PERMISSION_DENIED("Export failed: " + status.ToString()); + } + THROW_IO_EXCEPTION("Export failed: " + status.ToString()); } - THROW_IO_EXCEPTION("Export failed: " + status.ToString()); } ctx.clear(); return ctx; diff --git a/src/compiler/function/json_export_function.cc b/src/compiler/function/json_export_function.cc index 55cf670c0..a220ca78a 100644 --- a/src/compiler/function/json_export_function.cc +++ b/src/compiler/function/json_export_function.cc @@ -19,34 +19,22 @@ #include #include #include +#include "neug/common/types/data_chunk.h" +#include "neug/common/types/value.h" +#include "neug/execution/common/operators/retrieve/sink.h" #include "neug/utils/io/stream/output_stream.h" #include +#include "neug/common/types/property_types.h" #include "neug/compiler/function/read_function.h" #include "neug/compiler/main/metadata_registry.h" -#include "neug/generated/proto/response/response.pb.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/io/write/writer.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { namespace writer { -#define TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(CASE_ENUM, GETTER_METHOD, TYPE) \ - case neug::Array::TypedArrayCase::CASE_ENUM: { \ - auto& typed_array = arr.GETTER_METHOD(); \ - if (!StringFormatBuffer::validateProtoValue(typed_array.validity(), \ - rowIdx)) { \ - RETURN_STATUS_ERROR( \ - neug::StatusCode::ERR_INVALID_ARGUMENT, \ - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); \ - } \ - rapidjson::Value v(static_cast(typed_array.values(rowIdx))); \ - return v; \ - } - static neug::result parseJsonStringToValue( const std::string& json_str, int rowIdx, rapidjson::Document& parse_doc, const char* type_name) { @@ -71,275 +59,156 @@ static neug::result parseJsonStringToValue( return v; } -// return `rapidjson::Value` directly will not lead to any memory allocation, -// it's a move operation -static neug::result formatValueToJson( - const neug::Array& arr, int rowIdx, rapidjson::Document& doc) { +static neug::result valueToJsonValue( + const Value& value, rapidjson::Document& doc) { auto& allocator = doc.GetAllocator(); - switch (arr.typed_array_case()) { - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kBoolArray, bool_array, bool) - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kInt32Array, int32_array, int32_t) - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kInt64Array, int64_array, int64_t) - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kUint32Array, uint32_array, uint32_t) - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kUint64Array, uint64_array, uint64_t) - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kFloatArray, float_array, float) - TYPED_PRIMITIVE_ARRAY_TO_JSON_VALUE(kDoubleArray, double_array, double) - case neug::Array::TypedArrayCase::kStringArray: { - auto& string_array = arr.string_array(); - if (!StringFormatBuffer::validateProtoValue(string_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& str = string_array.values(rowIdx); + if (value.IsNull()) { + return rapidjson::Value(rapidjson::kNullType); + } + switch (value.type().id()) { + case DataTypeId::kBoolean: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kInt32: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kInt64: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kUInt32: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kUInt64: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kFloat: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kDouble: + return rapidjson::Value(value.GetValue()); + case DataTypeId::kVarchar: { + const auto& str = StringValue::Get(value); rapidjson::Value v; v.SetString(str.c_str(), static_cast(str.size()), allocator); return v; } - case neug::Array::TypedArrayCase::kDateArray: { - auto& date32_arr = arr.date_array(); - if (!StringFormatBuffer::validateProtoValue(date32_arr.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - Date date_value; - date_value.from_timestamp(date32_arr.values(rowIdx)); - const auto& s = date_value.to_string(); + case DataTypeId::kDate: { + const auto& s = value.GetValue().to_string(); rapidjson::Value v; v.SetString(s.c_str(), static_cast(s.size()), allocator); return v; } - case neug::Array::TypedArrayCase::kTimestampArray: { - auto& timestamp_array = arr.timestamp_array(); - if (!StringFormatBuffer::validateProtoValue(timestamp_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - DateTime dt_value(timestamp_array.values(rowIdx)); - const auto& s = dt_value.to_string(); + case DataTypeId::kTimestampMs: { + const auto& s = value.GetValue().to_string(); rapidjson::Value v; v.SetString(s.c_str(), static_cast(s.size()), allocator); return v; } - case neug::Array::TypedArrayCase::kIntervalArray: { - auto& interval_array = arr.interval_array(); - if (!StringFormatBuffer::validateProtoValue(interval_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& s = interval_array.values(rowIdx); + case DataTypeId::kInterval: { + const auto& s = value.GetValue().to_string(); rapidjson::Value v; v.SetString(s.c_str(), static_cast(s.size()), allocator); return v; } - case neug::Array::TypedArrayCase::kListArray: { - auto& list_array = arr.list_array(); - if (!StringFormatBuffer::validateProtoValue(list_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - rapidjson::Value arr_val(rapidjson::kArrayType); - uint32_t list_size = - list_array.offsets(rowIdx + 1) - list_array.offsets(rowIdx); - size_t offset = list_array.offsets(rowIdx); - for (uint32_t i = 0; i < list_size; ++i) { - rapidjson::Value elem; - GS_ASSIGN(elem, formatValueToJson(list_array.elements(), - static_cast(offset + i), doc)); - arr_val.PushBack(std::move(elem), allocator); - } - return arr_val; - } - case neug::Array::TypedArrayCase::kStructArray: { - auto& struct_arr = arr.struct_array(); - if (!StringFormatBuffer::validateProtoValue(struct_arr.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - rapidjson::Value arr_val(rapidjson::kArrayType); - for (int i = 0; i < struct_arr.fields_size(); ++i) { - const auto& field = struct_arr.fields(i); - rapidjson::Value elem; - GS_ASSIGN(elem, formatValueToJson(field, rowIdx, doc)); - arr_val.PushBack(std::move(elem), allocator); - } - return arr_val; - } - case neug::Array::TypedArrayCase::kVertexArray: { - auto& vertex_array = arr.vertex_array(); - if (!StringFormatBuffer::validateProtoValue(vertex_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - return parseJsonStringToValue(vertex_array.values(rowIdx), rowIdx, doc, - "vertex"); - } - case neug::Array::TypedArrayCase::kEdgeArray: { - auto& edge_array = arr.edge_array(); - if (!StringFormatBuffer::validateProtoValue(edge_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - return parseJsonStringToValue(edge_array.values(rowIdx), rowIdx, doc, - "edge"); - } - case neug::Array::TypedArrayCase::kPathArray: { - auto& path_array = arr.path_array(); - if (!StringFormatBuffer::validateProtoValue(path_array.validity(), - rowIdx)) { - RETURN_STATUS_ERROR(neug::StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - return parseJsonStringToValue(path_array.values(rowIdx), rowIdx, doc, - "path"); - } - default: - RETURN_STATUS_ERROR( - neug::StatusCode::ERR_INVALID_ARGUMENT, - "Unsupported type: " + std::to_string(arr.typed_array_case())); + default: { + const auto& s = value.to_string(); + rapidjson::Value v; + v.SetString(s.c_str(), static_cast(s.size()), + allocator); + return v; } -} - -static std::string getColumnName(const reader::EntrySchema& entry_schema, - size_t colIdx) { - if (colIdx < entry_schema.columnNames.size()) { - return entry_schema.columnNames[colIdx]; } - LOG(WARNING) << "Column index out of range: colIdx=" << colIdx - << ", using default column name"; - return "col_" + std::to_string(colIdx); } -JsonArrayStringFormatBuffer::JsonArrayStringFormatBuffer( - const neug::QueryResponse* response, const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema) - : StringFormatBuffer(response, schema), entry_schema_(entry_schema) { - buffer_.SetArray(); - current_line_.SetObject(); -} - -void JsonArrayStringFormatBuffer::addValue(int rowIdx, int colIdx) { - if (!validateIndex(response_, rowIdx, colIdx)) { - THROW_IO_EXCEPTION( - "Value index out of range: rowIdx=" + std::to_string(rowIdx) + - ", colIdx=" + std::to_string(colIdx)); - } - const neug::Array& column = response_->arrays(colIdx); - auto jsonResult = formatValueToJson(column, rowIdx, document_); - auto& allocator = document_.GetAllocator(); - WriteOptions writeOpts; - bool ignoreErrors = writeOpts.ignore_errors.get(schema_.options); - if (!jsonResult && !ignoreErrors) { - THROW_IO_EXCEPTION( - "Format value to JSON failed, rowIdx=" + std::to_string(rowIdx) + - ", colIdx=" + std::to_string(colIdx) + - ", error=" + jsonResult.error().ToString()); - } - const auto& columnName = getColumnName(entry_schema_, colIdx); - rapidjson::Value key(columnName.c_str(), - static_cast(columnName.size()), - allocator); - if (jsonResult) { - current_line_.AddMember(key, std::move(*jsonResult), allocator); - } else { - // add null value to ignore errors - current_line_.AddMember(key, rapidjson::Value(rapidjson::kNullType), - allocator); - } - if (colIdx == static_cast(response_->arrays_size()) - 1) { - buffer_.PushBack(std::move(current_line_), allocator); - current_line_.SetObject(); - } +static bool isSerializedGraphJson(DataTypeId id) { + return id == DataTypeId::kVertex || id == DataTypeId::kEdge || + id == DataTypeId::kPath; } -neug::Status JsonArrayStringFormatBuffer::flush(io::OutputStream& stream) { - if (buffer_.IsArray() && buffer_.Empty()) { - return neug::Status::OK(); +static const DataType& exportSourceType( + const std::vector& source_types, size_t col) { + static const DataType kDefaultVarchar(DataTypeId::kVarchar); + if (col < source_types.size()) { + return source_types[col]; } - const auto& jsonStr = rapidjson_stringify(buffer_); - buffer_.Clear(); - return stream.Write(reinterpret_cast(jsonStr.data()), - static_cast(jsonStr.size())); + return kDefaultVarchar; } -JsonLStringFormatBuffer::JsonLStringFormatBuffer( - const neug::QueryResponse* response, const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema) - : StringFormatBuffer(response, schema), entry_schema_(entry_schema) { - current_line_.SetObject(); - WriteOptions writeOpts; - size_t batchSize = writeOpts.batch_rows.get(schema.options); - if (batchSize > 0 && response->row_count() > 0) { - buffer_.reserve(batchSize); +// Convert one cell to a rapidjson value. When source_types[col] is a graph +// type, the VARCHAR payload is pre-serialized JSON and is parsed back into a +// structured value for inline emission. +static neug::result cellToJsonValue( + const IContextColumn& column, size_t row, const DataType& source_type, + rapidjson::Document& doc) { + Value value = column.get_elem(row); + if (isSerializedGraphJson(source_type.id())) { + const auto& str = StringValue::Get(value); + if (str.empty()) { + return rapidjson::Value(rapidjson::kNullType); + } + return parseJsonStringToValue(str, static_cast(row), doc, + source_type.ToString().c_str()); } + return valueToJsonValue(value, doc); } -void JsonLStringFormatBuffer::addValue(int rowIdx, int colIdx) { - if (!validateIndex(response_, rowIdx, colIdx)) { - THROW_IO_EXCEPTION( - "Value index out of range: rowIdx=" + std::to_string(rowIdx) + - ", colIdx=" + std::to_string(colIdx)); - } - const neug::Array& column = response_->arrays(colIdx); - auto jsonResult = formatValueToJson(column, rowIdx, document_); - auto& allocator = document_.GetAllocator(); - WriteOptions writeOpts; - bool ignoreErrors = writeOpts.ignore_errors.get(schema_.options); - if (!jsonResult && !ignoreErrors) { - THROW_IO_EXCEPTION( - "Format value to JSON failed, rowIdx=" + std::to_string(rowIdx) + - ", colIdx=" + std::to_string(colIdx) + - ", error=" + jsonResult.error().ToString()); - } - const auto& columnName = getColumnName(entry_schema_, colIdx); - rapidjson::Value key(columnName.c_str(), - static_cast(columnName.size()), - allocator); - if (jsonResult) { - current_line_.AddMember(key, std::move(*jsonResult), allocator); - } else { - current_line_.AddMember(key, rapidjson::Value(rapidjson::kNullType), - allocator); +static Status writeChunkAsJsonArray(const DataChunk& chunk, + const reader::FileSchema& schema, + const reader::EntrySchema& entry_schema, + const std::vector& source_types) { + if (schema.paths.empty()) { + return Status(StatusCode::ERR_INVALID_ARGUMENT, "Schema paths is empty"); } - if (colIdx == static_cast(response_->arrays_size()) - 1) { - buffer_.push_back(std::move(current_line_)); - current_line_.SetObject(); + auto stream = io::openLocalOutputStream(schema.paths[0]); + if (!stream) { + return Status(StatusCode::ERR_IO_ERROR, "Failed to open output file"); } -} -neug::Status JsonLStringFormatBuffer::flush(io::OutputStream& stream) { - for (const auto& val : buffer_) { - const auto& jsonStr = rapidjson_stringify(val); - auto status = stream.Write(reinterpret_cast(jsonStr.data()), - static_cast(jsonStr.size())); - if (!status.ok()) { - return status; - } - status = stream.Write( - reinterpret_cast(DEFAULT_JSON_NEWLINE), sizeof(char)); - if (!status.ok()) { - return status; + rapidjson::Document doc; + doc.SetArray(); + auto& allocator = doc.GetAllocator(); + for (size_t row = 0; row < chunk.row_num(); ++row) { + rapidjson::Value line(rapidjson::kObjectType); + for (size_t col = 0; col < chunk.col_num(); ++col) { + if (chunk.columns[col] == nullptr) { + continue; + } + const auto& column_name = col < entry_schema.columnNames.size() + ? entry_schema.columnNames[col] + : ("col_" + std::to_string(col)); + rapidjson::Value key(column_name.c_str(), + static_cast(column_name.size()), + allocator); + if (!chunk.columns[col]->has_value(row)) { + line.AddMember(key, rapidjson::Value(rapidjson::kNullType), allocator); + continue; + } + auto json_val = cellToJsonValue(*chunk.columns[col], row, + exportSourceType(source_types, col), doc); + if (!json_val) { + (void) stream->Close(); + return json_val.error(); + } + line.AddMember(key, std::move(*json_val), allocator); } + doc.PushBack(line, allocator); } - buffer_.clear(); - return neug::Status::OK(); + + rapidjson::StringBuffer buffer; + rapidjson::Writer writer(buffer); + doc.Accept(writer); + auto status = + stream->Write(reinterpret_cast(buffer.GetString()), + static_cast(buffer.GetSize())); + if (!status.ok()) { + (void) stream->Close(); + return status; + } + return stream->Close(); } -static Status writeTableWithBuffer(StringFormatBuffer& buffer, - const reader::FileSchema& schema, - const neug::QueryResponse* table, - size_t batchSize) { +static Status writeChunkAsJsonL(const DataChunk& chunk, + const reader::FileSchema& schema, + const reader::EntrySchema& entry_schema, + const std::vector& source_types) { if (schema.paths.empty()) { return Status(StatusCode::ERR_INVALID_ARGUMENT, "Schema paths is empty"); } @@ -348,55 +217,80 @@ static Status writeTableWithBuffer(StringFormatBuffer& buffer, return Status(StatusCode::ERR_IO_ERROR, "Failed to open output file"); } - if (batchSize == 0) { - return Status(StatusCode::ERR_INVALID_ARGUMENT, - "Batch size should be positive"); - } - - for (size_t i = 0; i < table->row_count(); ++i) { - for (size_t j = 0; j < table->arrays_size(); ++j) { - buffer.addValue(static_cast(i), static_cast(j)); - } - if ((i + 1) % static_cast(batchSize) == 0) { - auto status = buffer.flush(*stream); - if (!status.ok()) { + for (size_t row = 0; row < chunk.row_num(); ++row) { + rapidjson::Document doc; + doc.SetObject(); + auto& allocator = doc.GetAllocator(); + for (size_t col = 0; col < chunk.col_num(); ++col) { + if (chunk.columns[col] == nullptr) { + continue; + } + const auto& column_name = col < entry_schema.columnNames.size() + ? entry_schema.columnNames[col] + : ("col_" + std::to_string(col)); + rapidjson::Value key(column_name.c_str(), + static_cast(column_name.size()), + allocator); + if (!chunk.columns[col]->has_value(row)) { + doc.AddMember(key, rapidjson::Value(rapidjson::kNullType), allocator); + continue; + } + auto json_val = cellToJsonValue(*chunk.columns[col], row, + exportSourceType(source_types, col), doc); + if (!json_val) { (void) stream->Close(); - return Status(StatusCode::ERR_IO_ERROR, - "Failed to flush JSON buffer: " + status.ToString()); + return json_val.error(); } + doc.AddMember(key, std::move(*json_val), allocator); + } + rapidjson::StringBuffer buffer; + rapidjson::Writer writer(buffer); + doc.Accept(writer); + auto status = + stream->Write(reinterpret_cast(buffer.GetString()), + static_cast(buffer.GetSize())); + if (!status.ok()) { + (void) stream->Close(); + return status; + } + status = stream->Write( + reinterpret_cast(DEFAULT_JSON_NEWLINE), sizeof(char)); + if (!status.ok()) { + (void) stream->Close(); + return status; } - } - - auto status = buffer.flush(*stream); - if (!status.ok()) { - (void) stream->Close(); - return Status(StatusCode::ERR_IO_ERROR, - "Failed to flush JSON buffer: " + status.ToString()); } return stream->Close(); } -Status JsonArrayExportWriter::writeTable(const neug::QueryResponse* table) { +Status JsonArrayExportWriter::write(const DataChunk& chunk, + const std::vector& source_types) { if (!entry_schema_) { return Status(StatusCode::ERR_INVALID_ARGUMENT, "entry_schema is null"); } - JsonArrayStringFormatBuffer buffer(table, schema_, *entry_schema_); - size_t batchSize = table->row_count(); - if (batchSize == 0) { - batchSize = 1; + if (!source_types.empty() && source_types.size() != chunk.col_num()) { + return Status(StatusCode::ERR_INVALID_ARGUMENT, + "source_types size mismatch: expected " + + std::to_string(chunk.col_num()) + ", got " + + std::to_string(source_types.size())); } - return writeTableWithBuffer(buffer, schema_, table, batchSize); + return writeChunkAsJsonArray(chunk, schema_, *entry_schema_, source_types); } -Status JsonLExportWriter::writeTable(const neug::QueryResponse* table) { +Status JsonLExportWriter::write(const DataChunk& chunk, + const std::vector& source_types) { if (!entry_schema_) { return Status(StatusCode::ERR_INVALID_ARGUMENT, "entry_schema is null"); } - JsonLStringFormatBuffer buffer(table, schema_, *entry_schema_); - WriteOptions writeOpts; - size_t batchSize = writeOpts.batch_rows.get(schema_.options); - return writeTableWithBuffer(buffer, schema_, table, batchSize); + if (!source_types.empty() && source_types.size() != chunk.col_num()) { + return Status(StatusCode::ERR_INVALID_ARGUMENT, + "source_types size mismatch: expected " + + std::to_string(chunk.col_num()) + ", got " + + std::to_string(source_types.size())); + } + return writeChunkAsJsonL(chunk, schema_, *entry_schema_, source_types); } + } // namespace writer namespace function { @@ -410,9 +304,13 @@ static execution::Context jsonExecFunc( } auto writer = std::make_shared( schema, entry_schema); - auto status = writer->write(ctx, graph); - if (!status.ok()) { - THROW_IO_EXCEPTION("Export failed: " + status.ToString()); + auto source_types = ctx.column_types(); + auto chunks = neug::execution::Sink::materialize_for_export(ctx, graph); + for (const auto& chunk : chunks) { + auto status = writer->write(chunk, source_types); + if (!status.ok()) { + THROW_IO_EXCEPTION("Export failed: " + status.ToString()); + } } ctx.clear(); return ctx; @@ -445,9 +343,13 @@ static execution::Context jsonLExecFunc( } auto writer = std::make_shared(schema, entry_schema); - auto status = writer->write(ctx, graph); - if (!status.ok()) { - THROW_IO_EXCEPTION("Export failed: " + status.ToString()); + auto source_types = ctx.column_types(); + auto chunks = neug::execution::Sink::materialize_for_export(ctx, graph); + for (const auto& chunk : chunks) { + auto status = writer->write(chunk, source_types); + if (!status.ok()) { + THROW_IO_EXCEPTION("Export failed: " + status.ToString()); + } } ctx.clear(); return ctx; diff --git a/src/execution/common/context.cc b/src/execution/common/context.cc index 97808a927..825fce195 100644 --- a/src/execution/common/context.cc +++ b/src/execution/common/context.cc @@ -54,6 +54,13 @@ void Context::append_chunk(ContextChunk&& chunk) { chunks_.push_back(std::move(chunk)); } +void Context::append_chunks(std::vector&& chunks) { + chunks_.reserve(chunks_.size() + chunks.size()); + for (auto& chunk : chunks) { + append_chunk(std::move(chunk)); + } +} + void Context::flatten() { if (chunks_.size() <= 1) { return; @@ -86,6 +93,21 @@ size_t Context::row_num() const { return total; } +std::vector Context::column_types() const { + std::vector types; + types.reserve(tag_ids.size()); + for (size_t tag_id : tag_ids) { + for (size_t c = 0; c < chunks_.size(); ++c) { + auto col = chunks_[c].get(tag_id); + if (col != nullptr) { + types.push_back(col->elem_type()); + break; + } + } + } + return types; +} + } // namespace execution } // namespace neug diff --git a/src/execution/common/operators/retrieve/join.cc b/src/execution/common/operators/retrieve/join.cc index fed615477..ffae7c01f 100644 --- a/src/execution/common/operators/retrieve/join.cc +++ b/src/execution/common/operators/retrieve/join.cc @@ -18,12 +18,12 @@ #include "neug/common/columns/vertex_columns.h" #include "neug/common/types.h" #include "neug/common/types/data_chunk.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context_chunk.h" #include "neug/execution/utils/params.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/encoder.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { diff --git a/src/execution/common/operators/retrieve/path_expand_impl.cc b/src/execution/common/operators/retrieve/path_expand_impl.cc index 8b95a3c42..688f15acb 100644 --- a/src/execution/common/operators/retrieve/path_expand_impl.cc +++ b/src/execution/common/operators/retrieve/path_expand_impl.cc @@ -15,7 +15,7 @@ #include "neug/execution/common/operators/retrieve/path_expand_impl.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/src/execution/common/operators/retrieve/sink.cc b/src/execution/common/operators/retrieve/sink.cc index ff9b326d6..2e02c8fc6 100644 --- a/src/execution/common/operators/retrieve/sink.cc +++ b/src/execution/common/operators/retrieve/sink.cc @@ -15,6 +15,7 @@ #include "neug/execution/common/operators/retrieve/sink.h" +#include "neug/common/columns/columns_utils.h" #include "neug/common/columns/edge_columns.h" #include "neug/common/columns/list_columns.h" #include "neug/common/columns/path_columns.h" @@ -22,13 +23,12 @@ #include "neug/common/columns/value_columns.h" #include "neug/common/columns/vertex_columns.h" #include "neug/common/types/array_columns.h" +#include "neug/common/types/data_chunk.h" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/execution/common/context.h" - #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" - #include "rapidjson/document.h" #include "rapidjson/stringbuffer.h" #include "rapidjson/writer.h" @@ -544,7 +544,6 @@ void Sink::sink_results(const Context& ctx, const StorageReadInterface& graph, response->mutable_arrays()->Reserve(ctx.tag_ids.size()); for (size_t i : ctx.tag_ids) { - // Merge column across all chunks via union_col. std::shared_ptr merged; for (size_t c = 0; c < ctx.chunk_num(); ++c) { auto col = ctx.chunk(c).get(i); @@ -563,5 +562,155 @@ void Sink::sink_results(const Context& ctx, const StorageReadInterface& graph, } } +namespace { + +std::shared_ptr materialize_graph_column_as_string( + const std::shared_ptr& col, + const StorageReadInterface& graph) { + auto builder = + std::make_shared>(col->is_optional()); + builder->reserve(col->size()); + + switch (col->elem_type().id()) { + case DataTypeId::kVertex: { + auto vcol = std::dynamic_pointer_cast(col); + for (size_t i = 0; i < col->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + builder->push_back_null(); + } else { + builder->push_back_elem(Value::CreateValue( + convert_vertex_to_json(graph, vcol->get_vertex(i)))); + } + } + break; + } + case DataTypeId::kEdge: { + auto ecol = std::dynamic_pointer_cast(col); + for (size_t i = 0; i < col->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + builder->push_back_null(); + } else { + builder->push_back_elem(Value::CreateValue( + convert_edge_to_json(graph, ecol->get_edge(i)))); + } + } + break; + } + case DataTypeId::kPath: { + auto pcol = std::dynamic_pointer_cast(col); + for (size_t i = 0; i < col->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + builder->push_back_null(); + } else { + builder->push_back_elem(Value::CreateValue( + convert_path_to_json(graph, pcol->get_path(i)))); + } + } + break; + } + default: + return col; + } + return builder->finish(); +} + +Value materialize_value_for_export(const Value& value, + const StorageReadInterface& graph) { + if (value.IsNull()) { + return Value(value.type()); + } + switch (value.type().id()) { + case DataTypeId::kVertex: + return Value::STRING( + convert_vertex_to_json(graph, value.GetValue())); + case DataTypeId::kEdge: + return Value::STRING(convert_edge_to_json(graph, value.GetValue())); + case DataTypeId::kPath: + return Value::STRING(convert_path_to_json(graph, PathValue::Get(value))); + case DataTypeId::kList: { + const auto& children = ListValue::GetChildren(value); + const auto& child_type = ListType::GetChildType(value.type()); + std::vector materialized; + materialized.reserve(children.size()); + for (const auto& child : children) { + materialized.push_back(materialize_value_for_export(child, graph)); + } + return Value::LIST(child_type, std::move(materialized)); + } + case DataTypeId::kStruct: { + const auto& children = StructValue::GetChildren(value); + std::vector materialized; + materialized.reserve(children.size()); + for (const auto& child : children) { + materialized.push_back(materialize_value_for_export(child, graph)); + } + return Value::STRUCT(value.type(), std::move(materialized)); + } + case DataTypeId::kArray: { + const auto& children = ArrayValue::GetChildren(value); + std::vector materialized; + materialized.reserve(children.size()); + for (const auto& child : children) { + materialized.push_back(materialize_value_for_export(child, graph)); + } + return Value::ARRAY(value.type(), std::move(materialized)); + } + default: + return value; + } +} + +std::shared_ptr materialize_column_for_export( + const std::shared_ptr& col, + const StorageReadInterface& graph) { + if (col == nullptr) { + return nullptr; + } + switch (col->elem_type().id()) { + case DataTypeId::kVertex: + case DataTypeId::kEdge: + case DataTypeId::kPath: + return materialize_graph_column_as_string(col, graph); + case DataTypeId::kList: + case DataTypeId::kStruct: + case DataTypeId::kArray: { + auto builder = ColumnsUtils::create_builder(col->elem_type()); + for (size_t i = 0; i < col->size(); ++i) { + if (col->is_optional() && !col->has_value(i)) { + builder->push_back_null(); + } else { + builder->push_back_elem( + materialize_value_for_export(col->get_elem(i), graph)); + } + } + return builder->finish(); + } + default: + return col; + } +} + +} // namespace + +std::vector Sink::materialize_for_export( + const Context& ctx, const StorageReadInterface& graph) { + std::vector result; + result.reserve(ctx.chunk_num()); + for (size_t c = 0; c < ctx.chunk_num(); ++c) { + const auto& ctx_chunk = ctx.chunk(c); + DataChunk chunk; + int alias = 0; + for (size_t tag_id : ctx.tag_ids) { + auto col = ctx_chunk.get(tag_id); + if (col == nullptr) { + continue; + } + chunk.set(alias++, materialize_column_for_export(col, graph)); + } + result.push_back(std::move(chunk)); + } + return result; +} + } // namespace execution } // namespace neug diff --git a/src/execution/execute/ops/retrieve/edge.cc b/src/execution/execute/ops/retrieve/edge.cc index 358264e36..3b8da08c9 100644 --- a/src/execution/execute/ops/retrieve/edge.cc +++ b/src/execution/execute/ops/retrieve/edge.cc @@ -15,11 +15,11 @@ #include "neug/execution/execute/ops/retrieve/edge.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/operators/retrieve/edge_expand.h" #include "neug/execution/expression/expr.h" #include "neug/execution/expression/predicates.h" #include "neug/execution/utils/pb_parse_utils.h" -#include "neug/utils/property/types.h" namespace neug { class Schema; diff --git a/src/execution/execute/ops/retrieve/group_by.cc b/src/execution/execute/ops/retrieve/group_by.cc index e8f1f61ea..f1971b3a0 100644 --- a/src/execution/execute/ops/retrieve/group_by.cc +++ b/src/execution/execute/ops/retrieve/group_by.cc @@ -15,13 +15,13 @@ #include "neug/execution/execute/ops/retrieve/group_by.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/execution/common/operators/retrieve/group_by.h" #include "neug/execution/common/operators/retrieve/project.h" #include "neug/execution/execute/ops/retrieve/group_by_utils.h" #include "neug/storages/graph/graph_interface.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/src/execution/execute/ops/retrieve/scan.cc b/src/execution/execute/ops/retrieve/scan.cc index dbd6924e7..0ec4a379f 100644 --- a/src/execution/execute/ops/retrieve/scan.cc +++ b/src/execution/execute/ops/retrieve/scan.cc @@ -17,11 +17,11 @@ #include "neug/common/columns/value_columns.h" #include "neug/common/columns/vertex_columns.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/operators/retrieve/scan.h" #include "neug/execution/execute/ops/retrieve/scan_utils.h" #include "neug/execution/expression/predicates.h" #include "neug/execution/utils/params.h" -#include "neug/utils/property/types.h" namespace neug { namespace execution { diff --git a/src/execution/execute/ops/retrieve/select.cc b/src/execution/execute/ops/retrieve/select.cc index 200d9e5a8..433c5dbaa 100644 --- a/src/execution/execute/ops/retrieve/select.cc +++ b/src/execution/execute/ops/retrieve/select.cc @@ -14,11 +14,11 @@ */ #include "neug/execution/execute/ops/retrieve/select.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/execution/common/operators/retrieve/select.h" #include "neug/execution/expression/special_predicates.h" #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" #include "neug/common/columns/vertex_columns.h" #include "neug/execution/expression/predicates.h" diff --git a/src/execution/execute/ops/retrieve/tc_fuse.cc b/src/execution/execute/ops/retrieve/tc_fuse.cc index fd1cb6be6..373169d29 100644 --- a/src/execution/execute/ops/retrieve/tc_fuse.cc +++ b/src/execution/execute/ops/retrieve/tc_fuse.cc @@ -14,6 +14,7 @@ */ #include "neug/common/types/graph_types.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/execution/common/operators/retrieve/edge_expand.h" #include "neug/execution/execute/operator.h" @@ -21,7 +22,6 @@ #include "neug/execution/expression/special_predicates.h" #include "neug/execution/utils/pb_parse_utils.h" #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" namespace neug { namespace execution { diff --git a/src/execution/execute/ops/retrieve/vertex.cc b/src/execution/execute/ops/retrieve/vertex.cc index a763fa7b1..4e0e7a34a 100644 --- a/src/execution/execute/ops/retrieve/vertex.cc +++ b/src/execution/execute/ops/retrieve/vertex.cc @@ -16,13 +16,13 @@ #include "neug/execution/execute/ops/retrieve/vertex.h" #include "neug/common/types/graph_types.h" +#include "neug/common/types/property_types.h" #include "neug/execution/common/context.h" #include "neug/execution/common/operators/retrieve/get_v.h" #include "neug/execution/expression/expr.h" #include "neug/execution/utils/params.h" #include "neug/execution/utils/pb_parse_utils.h" #include "neug/storages/graph/graph_interface.h" -#include "neug/utils/property/types.h" namespace neug { class Schema; diff --git a/src/execution/utils/pb_parse_utils.cc b/src/execution/utils/pb_parse_utils.cc index 7fe66fc0c..b29878b48 100644 --- a/src/execution/utils/pb_parse_utils.cc +++ b/src/execution/utils/pb_parse_utils.cc @@ -25,12 +25,12 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/generated/proto/plan/algebra.pb.h" #include "neug/generated/proto/plan/common.pb.h" #include "neug/generated/proto/plan/physical.pb.h" #include "neug/generated/proto/plan/type.pb.h" #include "neug/storages/graph/schema.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/src/main/session.cc b/src/main/session.cc index 645b046be..57a5bd892 100644 --- a/src/main/session.cc +++ b/src/main/session.cc @@ -30,6 +30,7 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/config.h" #include "neug/execution/common/operators/retrieve/sink.h" #include "neug/execution/common/params_map.h" @@ -54,7 +55,6 @@ #include "neug/utils/encoder.h" #include "neug/utils/likely.h" #include "neug/utils/pb_utils.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { diff --git a/src/storages/CMakeLists.txt b/src/storages/CMakeLists.txt index dbb667107..bf6ddf5a0 100644 --- a/src/storages/CMakeLists.txt +++ b/src/storages/CMakeLists.txt @@ -3,6 +3,7 @@ add_subdirectory(csr) add_subdirectory(loader) add_subdirectory(container) add_subdirectory(module) -set(NEUG_STORAGES_OBJFILES $ $ $ $ $) +add_subdirectory(property) +set(NEUG_STORAGES_OBJFILES $ $ $ $ $ $) set(ALL_OBJECT_FILES ${ALL_OBJECT_FILES} ${NEUG_STORAGES_OBJFILES} PARENT_SCOPE) diff --git a/src/storages/csr/csr_view_utils.cc b/src/storages/csr/csr_view_utils.cc index 5d52ba844..45be605c9 100644 --- a/src/storages/csr/csr_view_utils.cc +++ b/src/storages/csr/csr_view_utils.cc @@ -15,7 +15,7 @@ #include "neug/storages/csr/csr_view_utils.h" #include "neug/common/types/graph_types.h" -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/src/storages/csr/immutable_csr.cc b/src/storages/csr/immutable_csr.cc index f4813efc1..237c36ab2 100644 --- a/src/storages/csr/immutable_csr.cc +++ b/src/storages/csr/immutable_csr.cc @@ -31,9 +31,9 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/container/container_utils.h" #include "neug/storages/container/i_container.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/src/storages/csr/mutable_csr.cc b/src/storages/csr/mutable_csr.cc index 5152fcf95..6e84c67ba 100644 --- a/src/storages/csr/mutable_csr.cc +++ b/src/storages/csr/mutable_csr.cc @@ -33,11 +33,11 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/container/container_utils.h" #include "neug/storages/container/file_mmap_container.h" #include "neug/utils/exception/exception.h" #include "neug/utils/io/file/file_utils.h" -#include "neug/utils/property/types.h" #include "neug/utils/spinlock.h" namespace neug { diff --git a/src/storages/graph/edge_table.cc b/src/storages/graph/edge_table.cc index a4bc15081..d1a006a86 100644 --- a/src/storages/graph/edge_table.cc +++ b/src/storages/graph/edge_table.cc @@ -30,6 +30,7 @@ #include #include "neug/common/columns/value_columns.h" +#include "neug/common/types/property_types.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/csr/csr_view_utils.h" #include "neug/storages/csr/immutable_csr.h" @@ -38,7 +39,6 @@ #include "neug/storages/module/type_name.h" #include "neug/storages/module_descriptor.h" #include "neug/utils/io/file/file_utils.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/src/storages/graph/property_graph.cc b/src/storages/graph/property_graph.cc index 77601233a..a8d6d2982 100644 --- a/src/storages/graph/property_graph.cc +++ b/src/storages/graph/property_graph.cc @@ -26,15 +26,15 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/graph/schema.h" +#include "neug/storages/indexers.h" #include "neug/storages/module/module_broker.h" +#include "neug/storages/property/column.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/indexers.h" #include "neug/utils/io/file/file_utils.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/types.h" #include "neug/utils/yaml_utils.h" namespace neug { diff --git a/src/storages/graph/schema.cc b/src/storages/graph/schema.cc index a709dacce..cf5a6e395 100644 --- a/src/storages/graph/schema.cc +++ b/src/storages/graph/schema.cc @@ -29,11 +29,11 @@ #include #include #include "neug/common/extra_type_info.h" +#include "neug/common/types/property_types.h" +#include "neug/storages/id_indexer.h" #include "neug/storages/module/module_factory.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/id_indexer.h" #include "neug/utils/pb_utils.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" diff --git a/src/storages/loader/loader_utils.cc b/src/storages/loader/loader_utils.cc index 6636f6105..b4c8e9080 100644 --- a/src/storages/loader/loader_utils.cc +++ b/src/storages/loader/loader_utils.cc @@ -43,9 +43,9 @@ #include "neug/common/columns/columns_utils.h" #include "neug/common/columns/value_columns.h" #include "neug/common/types/value.h" +#include "neug/storages/property/column.h" #include "neug/utils/datetime_parsers.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/column.h" #include "neug/utils/string_utils.h" namespace neug { diff --git a/src/storages/property/CMakeLists.txt b/src/storages/property/CMakeLists.txt new file mode 100644 index 000000000..f8025cc00 --- /dev/null +++ b/src/storages/property/CMakeLists.txt @@ -0,0 +1,6 @@ +file(GLOB_RECURSE PROPERTY_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/*.cc") +add_library(neug_storages_property OBJECT ${PROPERTY_SOURCES}) +set(ALL_OBJECT_FILES + ${ALL_OBJECT_FILES} $ + PARENT_SCOPE) +add_dependencies(neug_storages_property neug_proto) diff --git a/src/utils/property/array_column.cc b/src/storages/property/array_column.cc similarity index 98% rename from src/utils/property/array_column.cc rename to src/storages/property/array_column.cc index ebe6b629b..067a97132 100644 --- a/src/utils/property/array_column.cc +++ b/src/storages/property/array_column.cc @@ -13,14 +13,14 @@ * limitations under the License. */ -#include "neug/utils/property/array_column.h" +#include "neug/storages/property/array_column.h" #include +#include "neug/common/types/property_types.h" #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/module/module_factory.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include diff --git a/src/utils/property/column.cc b/src/storages/property/column.cc similarity index 94% rename from src/utils/property/column.cc rename to src/storages/property/column.cc index 0000c1a9e..30cdc0a0a 100644 --- a/src/utils/property/column.cc +++ b/src/storages/property/column.cc @@ -13,17 +13,17 @@ * limitations under the License. */ -#include "neug/utils/property/column.h" +#include "neug/storages/property/column.h" #include +#include "neug/common/types/property_types.h" #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/container/container_utils.h" +#include "neug/storages/id_indexer.h" #include "neug/storages/module/module_factory.h" -#include "neug/utils/id_indexer.h" -#include "neug/utils/property/array_column.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" +#include "neug/storages/property/array_column.h" +#include "neug/storages/property/table.h" #include "neug/utils/serialization/out_archive.h" namespace neug { diff --git a/src/utils/property/table.cc b/src/storages/property/table.cc similarity index 99% rename from src/utils/property/table.cc rename to src/storages/property/table.cc index 58e6143ea..b5c8244c6 100644 --- a/src/utils/property/table.cc +++ b/src/storages/property/table.cc @@ -13,7 +13,7 @@ * limitations under the License. */ -#include "neug/utils/property/table.h" +#include "neug/storages/property/table.h" #include #include @@ -21,8 +21,8 @@ #include #include +#include "neug/storages/property/column.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/column.h" #include "neug/utils/serialization/out_archive.h" namespace neug { diff --git a/src/transaction/insert_transaction.cc b/src/transaction/insert_transaction.cc index 2ef0cf57b..b0f2334f9 100644 --- a/src/transaction/insert_transaction.cc +++ b/src/transaction/insert_transaction.cc @@ -22,13 +22,13 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/allocators.h" #include "neug/storages/graph/schema.h" #include "neug/transaction/transaction_utils.h" #include "neug/transaction/version_manager.h" #include "neug/transaction/wal/wal.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/out_archive.h" namespace neug { diff --git a/src/transaction/update_transaction.cc b/src/transaction/update_transaction.cc index 15b306608..c6cf67d79 100644 --- a/src/transaction/update_transaction.cc +++ b/src/transaction/update_transaction.cc @@ -26,22 +26,22 @@ #include #include "neug/common/extra_type_info.h" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/allocators.h" #include "neug/storages/csr/csr_base.h" #include "neug/storages/csr/csr_view_utils.h" #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" +#include "neug/storages/id_indexer.h" +#include "neug/storages/property/column.h" +#include "neug/storages/property/table.h" #include "neug/transaction/transaction_utils.h" #include "neug/transaction/version_manager.h" #include "neug/transaction/wal/wal.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/id_indexer.h" #include "neug/utils/io/file/file_utils.h" #include "neug/utils/likely.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" #include "neug/utils/serialization/out_archive.h" diff --git a/src/utils/bolt_utils.cc b/src/utils/bolt_utils.cc index 6be23908f..0f0f61a7d 100644 --- a/src/utils/bolt_utils.cc +++ b/src/utils/bolt_utils.cc @@ -31,8 +31,8 @@ #include #include "glog/logging.h" +#include "neug/common/types/property_types.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/src/utils/io/read/csv/csv_reader.cc b/src/utils/io/read/csv/csv_reader.cc index f7f7fc23e..120f3d500 100644 --- a/src/utils/io/read/csv/csv_reader.cc +++ b/src/utils/io/read/csv/csv_reader.cc @@ -35,7 +35,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" @@ -485,8 +484,8 @@ CsvReader::CsvReader(std::shared_ptr sharedState, CsvReader::~CsvReader() = default; -void CsvReader::read(std::shared_ptr /*localState*/, - execution::Context& ctx) { +std::vector CsvReader::read( + std::shared_ptr /*localState*/) { if (!sharedState_) { THROW_INVALID_ARGUMENT_EXCEPTION("SharedState is null"); } @@ -531,15 +530,14 @@ void CsvReader::read(std::shared_ptr /*localState*/, } if (use_batch_read && !sharedState_->skipRows) { - batch_read(suppliers, ctx); - } else { - full_read(suppliers, ctx, config); + return batch_read(suppliers); } + return full_read(suppliers, config); } -void CsvReader::full_read( +std::vector CsvReader::full_read( const std::vector>& suppliers, - execution::Context& output, const CsvReadConfig& output_config) { + const CsvReadConfig& output_config) { auto merged = read_all_chunks(suppliers); int expected_cols = sharedState_->columnNum(); @@ -559,19 +557,20 @@ void CsvReader::full_read( ? output_config.include_columns : sharedState_->projectColumns); - output.clear(); - output.append_chunk(std::move(projected)); + std::vector output; + output.push_back(std::move(projected)); + return output; } -void CsvReader::batch_read( - const std::vector>& suppliers, - execution::Context& output) { - output.clear(); +std::vector CsvReader::batch_read( + const std::vector>& suppliers) { + std::vector output; for (const auto& supplier : suppliers) { while (auto chunk = supplier->GetNextChunk()) { - output.append_chunk(std::move(*chunk)); + output.push_back(std::move(*chunk)); } } + return output; } result> CsvReader::inferSchema() { diff --git a/src/utils/io/read/json/json_reader.cc b/src/utils/io/read/json/json_reader.cc index fa7d0f933..025bad412 100644 --- a/src/utils/io/read/json/json_reader.cc +++ b/src/utils/io/read/json/json_reader.cc @@ -30,7 +30,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" @@ -297,8 +296,8 @@ JsonReader::JsonReader(std::shared_ptr sharedState, JsonReader::~JsonReader() = default; -void JsonReader::read(std::shared_ptr /*localState*/, - execution::Context& ctx) { +std::vector JsonReader::read( + std::shared_ptr /*localState*/) { if (!sharedState_ || !optionsBuilder_) { THROW_INVALID_ARGUMENT_EXCEPTION("JsonReader state or builder is null"); } @@ -340,15 +339,14 @@ void JsonReader::read(std::shared_ptr /*localState*/, } if (use_batch_read && !sharedState_->skipRows) { - batch_read(suppliers, ctx); - } else { - full_read(suppliers, ctx, config); + return batch_read(suppliers); } + return full_read(suppliers, config); } -void JsonReader::full_read( +std::vector JsonReader::full_read( const std::vector>& suppliers, - execution::Context& output, const JsonReadConfig& output_config) { + const JsonReadConfig& output_config) { auto merged = read_all_chunks(suppliers); int expected_cols = sharedState_->columnNum(); @@ -367,19 +365,20 @@ void JsonReader::full_read( sharedState_->projectColumns.empty() ? output_config.include_columns : sharedState_->projectColumns); - output.clear(); - output.append_chunk(std::move(projected)); + std::vector output; + output.push_back(std::move(projected)); + return output; } -void JsonReader::batch_read( - const std::vector>& suppliers, - execution::Context& output) { - output.clear(); +std::vector JsonReader::batch_read( + const std::vector>& suppliers) { + std::vector output; for (const auto& supplier : suppliers) { while (auto chunk = supplier->GetNextChunk()) { - output.append_chunk(std::move(*chunk)); + output.push_back(std::move(*chunk)); } } + return output; } result> JsonReader::inferSchema() { diff --git a/src/utils/io/write/writer.cc b/src/utils/io/write/writer.cc index d5aa5984e..8c3680698 100644 --- a/src/utils/io/write/writer.cc +++ b/src/utils/io/write/writer.cc @@ -14,12 +14,11 @@ */ #include "neug/utils/io/write/writer.h" -#include "neug/execution/common/operators/retrieve/sink.h" -#include "neug/generated/proto/response/response.pb.h" +#include "neug/common/types/property_types.h" +#include "neug/common/types/value.h" #include "neug/utils/exception/exception.h" #include "neug/utils/io/read/common/options.h" #include "neug/utils/io/stream/output_stream.h" -#include "neug/utils/property/types.h" #include #include @@ -31,71 +30,33 @@ namespace neug { namespace writer { -bool StringFormatBuffer::validateIndex(const neug::QueryResponse* response, - int rowIdx, int colIdx) { - if (response == nullptr) - return false; - if (rowIdx < 0 || rowIdx >= response->row_count()) - return false; - if (colIdx < 0 || static_cast(colIdx) >= response->arrays_size()) { - return false; - } - return true; -} - -bool StringFormatBuffer::validateProtoValue(const std::string& validity, - int rowIdx) { - return validity.empty() || - (static_cast(validity[static_cast(rowIdx) >> 3]) >> - (rowIdx & 7)) & - 1; -} - -#define TYPED_PRIMITIVE_ARRAY_TO_JSON(CASE_ENUM, GETTER_METHOD) \ - case neug::Array::TypedArrayCase::CASE_ENUM: { \ - auto& typed_array = arr.GETTER_METHOD(); \ - if (!validateProtoValue(typed_array.validity(), rowIdx)) { \ - return neug::Status( \ - StatusCode::ERR_INVALID_ARGUMENT, \ - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); \ - } \ - const auto& str = std::to_string(typed_array.values(rowIdx)); \ - write(reinterpret_cast(str.c_str()), str.size()); \ - return neug::Status::OK(); \ - } - -CSVStringFormatBuffer::CSVStringFormatBuffer( - const neug::QueryResponse* response, const reader::FileSchema& schema, +DataChunkCSVStringFormatBuffer::DataChunkCSVStringFormatBuffer( + const DataChunk& chunk, const reader::FileSchema& schema, const reader::EntrySchema& entry_schema) - : StringFormatBuffer(response, schema), entry_schema_(entry_schema) { - capacity_ = DEFAULT_CAPACITY; - WriteOptions writeOpts; - size_t batchSize = writeOpts.batch_rows.get(schema.options); - if (batchSize > 0 && response->arrays_size() > 0 && - response->row_count() > 0) { - size_t ncol = static_cast(response->arrays_size()); - if (batchSize <= SIZE_MAX / DEFAULT_CAPACITY && - ncol <= SIZE_MAX / (DEFAULT_CAPACITY * batchSize)) { - capacity_ = DEFAULT_CAPACITY * batchSize * ncol; - } else { - LOG(WARNING) << "CSV buffer capacity overflow, batchSize=" << batchSize - << ", ncol=" << ncol - << ", using default capacity: " << capacity_; - } - } - has_header_ = writeOpts.has_header.get(schema.options); - delimiter_ = writeOpts.delimiter.get(schema.options); - ignore_errors_ = writeOpts.ignore_errors.get(schema.options); - escape_char_ = writeOpts.escape_char.get(schema.options); - quote_char_ = writeOpts.quote_char.get(schema.options); + : chunk_(chunk), + schema_(schema), + entry_schema_(entry_schema), + capacity_(DEFAULT_CAPACITY) { + WriteOptions write_opts; + size_t batch_size = write_opts.batch_rows.get(schema.options); + if (batch_size > 0 && chunk.col_num() > 0 && chunk.row_num() > 0) { + size_t ncol = chunk.col_num(); + if (batch_size <= SIZE_MAX / DEFAULT_CAPACITY && + ncol <= SIZE_MAX / (DEFAULT_CAPACITY * batch_size)) { + capacity_ = DEFAULT_CAPACITY * batch_size * ncol; + } + } + has_header_ = write_opts.has_header.get(schema.options); + delimiter_ = write_opts.delimiter.get(schema.options); + ignore_errors_ = write_opts.ignore_errors.get(schema.options); + escape_char_ = write_opts.escape_char.get(schema.options); + quote_char_ = write_opts.quote_char.get(schema.options); blob_.data = std::make_unique(capacity_); blob_.size = 0; data_ = blob_.data.get(); } -void CSVStringFormatBuffer::addHeader() { - // Emit header at init so empty result sets still get a header row when - // HEADER = true. +void DataChunkCSVStringFormatBuffer::addHeader() { if (has_header_ && !entry_schema_.columnNames.empty()) { for (size_t col = 0; col < entry_schema_.columnNames.size(); ++col) { if (col > 0) { @@ -108,168 +69,34 @@ void CSVStringFormatBuffer::addHeader() { } } -neug::Status CSVStringFormatBuffer::formatValueToStr(const neug::Array& arr, - int rowIdx) { - switch (arr.typed_array_case()) { - TYPED_PRIMITIVE_ARRAY_TO_JSON(kBoolArray, bool_array) - TYPED_PRIMITIVE_ARRAY_TO_JSON(kInt32Array, int32_array) - TYPED_PRIMITIVE_ARRAY_TO_JSON(kInt64Array, int64_array) - TYPED_PRIMITIVE_ARRAY_TO_JSON(kUint32Array, uint32_array) - TYPED_PRIMITIVE_ARRAY_TO_JSON(kUint64Array, uint64_array) - TYPED_PRIMITIVE_ARRAY_TO_JSON(kFloatArray, float_array) - TYPED_PRIMITIVE_ARRAY_TO_JSON(kDoubleArray, double_array) - case neug::Array::TypedArrayCase::kStringArray: { - auto& string_array = arr.string_array(); - if (!validateProtoValue(string_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& str = string_array.values(rowIdx); - // add quotes for string type values - write(reinterpret_cast("e_char_), sizeof(char)); - // espace special characters - char escapeChars[] = {escape_char_, quote_char_}; - writeWithEscapes(escapeChars, escape_char_, str); - write(reinterpret_cast("e_char_), sizeof(char)); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kDateArray: { - auto& date32_arr = arr.date_array(); - if (!validateProtoValue(date32_arr.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - Date date_value; - date_value.from_timestamp(date32_arr.values(rowIdx)); - const auto& date_str = date_value.to_string(); - write(reinterpret_cast(date_str.c_str()), date_str.size()); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kTimestampArray: { - auto& timestamp_array = arr.timestamp_array(); - if (!validateProtoValue(timestamp_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - DateTime dt_value(timestamp_array.values(rowIdx)); - const auto& dt_str = dt_value.to_string(); - write(reinterpret_cast(dt_str.c_str()), dt_str.size()); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kIntervalArray: { - auto& interval_array = arr.interval_array(); - if (!validateProtoValue(interval_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& interval_str = interval_array.values(rowIdx); - write(reinterpret_cast(interval_str.c_str()), - interval_str.size()); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kListArray: { - auto& list_array = arr.list_array(); - if (!validateProtoValue(list_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - write(reinterpret_cast(&LIST_ARRAY_CHAR[0]), sizeof(char)); - uint32_t list_size = - list_array.offsets(rowIdx + 1) - list_array.offsets(rowIdx); - size_t offset = list_array.offsets(rowIdx); - for (uint32_t i = 0; i < list_size; ++i) { - if (i > 0) { - write(reinterpret_cast(COMMA_CHAR), sizeof(char)); - } - RETURN_IF_NOT_OK(formatValueToStr(list_array.elements(), offset + i)); - } - write(reinterpret_cast(&LIST_ARRAY_CHAR[1]), sizeof(char)); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kStructArray: { - auto& struct_arr = arr.struct_array(); - if (!validateProtoValue(struct_arr.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - write(reinterpret_cast(&LIST_ARRAY_CHAR[0]), sizeof(char)); - for (int i = 0; i < struct_arr.fields_size(); ++i) { - if (i > 0) { - write(reinterpret_cast(COMMA_CHAR), sizeof(char)); - } - const auto& field = struct_arr.fields(i); - RETURN_IF_NOT_OK(formatValueToStr(field, rowIdx)); - } - write(reinterpret_cast(&LIST_ARRAY_CHAR[1]), sizeof(char)); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kVertexArray: { - auto vertex_array = arr.vertex_array(); - if (!validateProtoValue(vertex_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& vertex_str = vertex_array.values(rowIdx); - write(reinterpret_cast(vertex_str.c_str()), - vertex_str.size()); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kEdgeArray: { - auto edge_array = arr.edge_array(); - if (!validateProtoValue(edge_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& edge_str = edge_array.values(rowIdx); - write(reinterpret_cast(edge_str.c_str()), edge_str.size()); - return neug::Status::OK(); - } - case neug::Array::TypedArrayCase::kPathArray: { - auto path_array = arr.path_array(); - if (!validateProtoValue(path_array.validity(), rowIdx)) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(rowIdx)); - } - const auto& path_str = path_array.values(rowIdx); - write(reinterpret_cast(path_str.c_str()), path_str.size()); - return neug::Status::OK(); - } - default: { - return neug::Status( - StatusCode::ERR_NOT_SUPPORTED, - "Unsupported type: " + std::to_string(arr.typed_array_case())); - } - } -} - -void CSVStringFormatBuffer::writeWithEscapes(char* toEscape, char escape, - const std::string& val) { +void DataChunkCSVStringFormatBuffer::writeWithEscapes(char* to_escape, + char escape, + const std::string& val) { uint64_t i = 0; - auto found = val.find_first_of(toEscape, 0, 2); - + auto found = val.find_first_of(to_escape, 0, 2); while (found != std::string::npos) { while (i < found) { write(reinterpret_cast(&val[i]), sizeof(char)); - i++; + ++i; } write(reinterpret_cast(&escape), sizeof(char)); - found = val.find_first_of(toEscape, found + sizeof(escape), 2); + found = val.find_first_of(to_escape, found + sizeof(escape), 2); } while (i < val.length()) { write(reinterpret_cast(&val[i]), sizeof(char)); - i++; + ++i; } } -void CSVStringFormatBuffer::write(const uint8_t* buffer, uint64_t len) { +void DataChunkCSVStringFormatBuffer::write(const uint8_t* buffer, + uint64_t len) { if (len == 0) { return; } if (buffer == nullptr) { - THROW_IO_EXCEPTION("CSVStringFormatBuffer::write called with null buffer"); + THROW_IO_EXCEPTION( + "DataChunkCSVStringFormatBuffer::write called with null buffer"); } - // Overflow-safe: need grow when (blob.size + len > capacity) without - // computing blob.size + len (which can overflow). const bool need_grow = (len > capacity_) || (blob_.size > capacity_ - len); if (need_grow) { size_t old_capacity = capacity_; @@ -280,7 +107,6 @@ void CSVStringFormatBuffer::write(const uint8_t* buffer, uint64_t len) { capacity_ *= 2; } while ((len > capacity_) || (blob_.size > capacity_ - len)); auto new_data = std::make_unique(capacity_); - // Copy only up to old capacity to avoid reading past old buffer size_t copy_len = (blob_.size < old_capacity) ? blob_.size : old_capacity; if (copy_len > 0 && data_ != nullptr) { memcpy(new_data.get(), data_, copy_len); @@ -289,40 +115,69 @@ void CSVStringFormatBuffer::write(const uint8_t* buffer, uint64_t len) { blob_.data = std::move(new_data); data_ = blob_.data.get(); } - memcpy(data_ + blob_.size, buffer, len); blob_.size += len; } -void CSVStringFormatBuffer::addValue(int rowIdx, int colIdx) { - if (!validateIndex(response_, rowIdx, colIdx)) { - THROW_IO_EXCEPTION( - "Value index out of range: rowIdx=" + std::to_string(rowIdx) + - ", colIdx=" + std::to_string(colIdx)); +neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( + const Value& value, size_t row_idx) { + if (value.IsNull()) { + return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, + "Value is invalid, rowIdx=" + std::to_string(row_idx)); + } + const auto type_id = value.type().id(); + if (type_id == DataTypeId::kVarchar) { + const auto& str = StringValue::Get(value); + write(reinterpret_cast("e_char_), sizeof(char)); + char escape_chars[] = {escape_char_, quote_char_}; + writeWithEscapes(escape_chars, escape_char_, str); + write(reinterpret_cast("e_char_), sizeof(char)); + return neug::Status::OK(); } - if (colIdx > 0) { + const auto& str = value.to_string(); + write(reinterpret_cast(str.c_str()), str.size()); + return neug::Status::OK(); +} + +void DataChunkCSVStringFormatBuffer::addValue(size_t row_idx, size_t col_idx) { + if (col_idx >= chunk_.col_num() || chunk_.columns[col_idx] == nullptr) { + THROW_IO_EXCEPTION("Column index out of range: colIdx=" + + std::to_string(col_idx)); + } + const auto& column = chunk_.columns[col_idx]; + if (row_idx >= column->size()) { + THROW_IO_EXCEPTION("Row index out of range: rowIdx=" + + std::to_string(row_idx)); + } + if (col_idx > 0) { write(reinterpret_cast(&delimiter_), sizeof(char)); } - const neug::Array& column = response_->arrays(colIdx); - auto strResult = formatValueToStr(column, rowIdx); - if (!strResult.ok()) { + if (!column->has_value(row_idx)) { if (!ignore_errors_) { - THROW_IO_EXCEPTION( - "Format value to string failed, rowIdx=" + std::to_string(rowIdx) + - ", colIdx=" + std::to_string(colIdx) + - ", error=" + strResult.ToString()); - } else { + THROW_IO_EXCEPTION("Value is invalid, rowIdx=" + std::to_string(row_idx) + + ", colIdx=" + std::to_string(col_idx)); + } + write(reinterpret_cast(DEFAULT_NULL_STR), + strlen(DEFAULT_NULL_STR)); + } else { + auto str_result = formatValueToStr(column->get_elem(row_idx), row_idx); + if (!str_result.ok()) { + if (!ignore_errors_) { + THROW_IO_EXCEPTION( + "Format value to string failed, rowIdx=" + std::to_string(row_idx) + + ", colIdx=" + std::to_string(col_idx) + + ", error=" + str_result.ToString()); + } write(reinterpret_cast(DEFAULT_NULL_STR), strlen(DEFAULT_NULL_STR)); } } - if (colIdx == response_->arrays_size() - 1) { - // the last column, add newline + if (col_idx + 1 == chunk_.col_num()) { write(reinterpret_cast(DEFAULT_CSV_NEWLINE), sizeof(char)); } } -neug::Status CSVStringFormatBuffer::flush(io::OutputStream& stream) { +neug::Status DataChunkCSVStringFormatBuffer::flush(io::OutputStream& stream) { if (blob_.size > 0) { auto status = stream.Write(data_, static_cast(blob_.size)); blob_.size = 0; @@ -333,15 +188,8 @@ neug::Status CSVStringFormatBuffer::flush(io::OutputStream& stream) { return neug::Status::OK(); } -neug::Status QueryExportWriter::write(const execution::Context& context, - const StorageReadInterface& graph) { - neug::QueryResponse response; - execution::Sink::sink_results(context, graph, &response); - return writeTable(&response); -} - -neug::Status CsvQueryExportWriter::writeTable( - const neug::QueryResponse* table) { +neug::Status CsvQueryExportWriter::write( + const DataChunk& chunk, const std::vector& /*source_types*/) { if (schema_.paths.empty()) { return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, "Schema paths is empty"); @@ -355,20 +203,21 @@ neug::Status CsvQueryExportWriter::writeTable( return neug::Status(StatusCode::ERR_IO_ERROR, "Failed to open output file"); } - WriteOptions writeOpts; - auto batchSize = writeOpts.batch_rows.get(schema_.options); - if (batchSize <= 0) { + WriteOptions write_opts; + auto batch_size = write_opts.batch_rows.get(schema_.options); + if (batch_size <= 0) { return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, "Batch size should be positive"); } - auto csvBuffer = CSVStringFormatBuffer(table, schema_, *entry_schema_); - csvBuffer.addHeader(); - for (size_t i = 0; i < table->row_count(); ++i) { - for (size_t j = 0; j < table->arrays_size(); ++j) { - csvBuffer.addValue(i, j); + auto csv_buffer = + DataChunkCSVStringFormatBuffer(chunk, schema_, *entry_schema_); + csv_buffer.addHeader(); + for (size_t i = 0; i < chunk.row_num(); ++i) { + for (size_t j = 0; j < chunk.col_num(); ++j) { + csv_buffer.addValue(i, j); } - if (i % batchSize == batchSize - 1) { - auto status = csvBuffer.flush(*stream); + if (i % batch_size == batch_size - 1) { + auto status = csv_buffer.flush(*stream); if (!status.ok()) { (void) stream->Close(); return neug::Status(StatusCode::ERR_IO_ERROR, @@ -377,7 +226,7 @@ neug::Status CsvQueryExportWriter::writeTable( } } - auto status = csvBuffer.flush(*stream); + auto status = csv_buffer.flush(*stream); if (!status.ok()) { (void) stream->Close(); return neug::Status(StatusCode::ERR_IO_ERROR, diff --git a/src/utils/pb_utils.cc b/src/utils/pb_utils.cc index 4063fb15c..5f0ca6caa 100644 --- a/src/utils/pb_utils.cc +++ b/src/utils/pb_utils.cc @@ -34,13 +34,13 @@ #include #include #include +#include "neug/common/types/default_value.h" +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/execution/expression/expr.h" #include "neug/generated/proto/plan/common.pb.h" #include "neug/generated/proto/plan/expr.pb.h" #include "neug/utils/bolt_utils.h" -#include "neug/utils/property/default_value.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { diff --git a/src/utils/yaml_utils.cc b/src/utils/yaml_utils.cc index 6d1a70b2e..4d6a6eb28 100644 --- a/src/utils/yaml_utils.cc +++ b/src/utils/yaml_utils.cc @@ -34,8 +34,8 @@ #include #include "neug/common/extra_type_info.h" +#include "neug/common/types/property_types.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" #include "neug/utils/service_utils.h" diff --git a/tests/storage/alter_property_test.cc b/tests/storage/alter_property_test.cc index a361a9513..ad92d930d 100644 --- a/tests/storage/alter_property_test.cc +++ b/tests/storage/alter_property_test.cc @@ -23,13 +23,13 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/csr/csr_base.h" #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" #include "neug/storages/loader/loader_utils.h" #include "neug/utils/io/read/csv/csv_read_config.h" -#include "neug/utils/property/types.h" #include "unittest/utils.h" namespace neug { diff --git a/tests/storage/test_vertex_table.cc b/tests/storage/test_vertex_table.cc index d6249ba6b..36f0456ff 100644 --- a/tests/storage/test_vertex_table.cc +++ b/tests/storage/test_vertex_table.cc @@ -21,13 +21,13 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/main/neug_db.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/graph/schema.h" #include "neug/storages/graph/vertex_table.h" #include "neug/transaction/transaction_utils.h" -#include "neug/utils/property/types.h" #include "unittest/utils.h" #include diff --git a/tests/unittest/logical_delete_test.cc b/tests/unittest/logical_delete_test.cc index be902966b..0a415c2cf 100644 --- a/tests/unittest/logical_delete_test.cc +++ b/tests/unittest/logical_delete_test.cc @@ -15,9 +15,9 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/storages/graph/property_graph.h" -#include "neug/utils/property/types.h" #include "neug/utils/yaml_utils.h" #include "unittest/utils.h" diff --git a/tests/unittest/schema_test.cc b/tests/unittest/schema_test.cc index c5a1faa59..ffbd54f00 100644 --- a/tests/unittest/schema_test.cc +++ b/tests/unittest/schema_test.cc @@ -25,8 +25,8 @@ #include "neug/common/types/value.h" +#include "neug/common/types/property_types.h" #include "neug/storages/graph/schema.h" -#include "neug/utils/property/types.h" #include "neug/utils/yaml_utils.h" using neug::DataType; diff --git a/tests/unittest/test_column.cc b/tests/unittest/test_column.cc index 03b154da6..504943f97 100644 --- a/tests/unittest/test_column.cc +++ b/tests/unittest/test_column.cc @@ -20,9 +20,9 @@ #include "neug/storages/checkpoint.h" #include "neug/storages/checkpoint_manager.h" +#include "neug/storages/property/array_column.h" +#include "neug/storages/property/column.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/array_column.h" -#include "neug/utils/property/column.h" #include "unittest/utils.h" namespace neug { diff --git a/tests/unittest/test_indexer.cc b/tests/unittest/test_indexer.cc index 71bca3918..f72e8f187 100644 --- a/tests/unittest/test_indexer.cc +++ b/tests/unittest/test_indexer.cc @@ -27,9 +27,9 @@ #include "neug/common/types/value.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/container/file_header.h" +#include "neug/storages/id_indexer.h" #include "neug/storages/module/module_factory.h" #include "neug/storages/module/type_name.h" -#include "neug/utils/id_indexer.h" #include "unittest/utils.h" namespace neug { diff --git a/tests/unittest/utils.h b/tests/unittest/utils.h index 8374ddbcd..1e7145ce5 100644 --- a/tests/unittest/utils.h +++ b/tests/unittest/utils.h @@ -27,19 +27,19 @@ #include "neug/common/columns/value_columns.h" #include "neug/common/types/data_chunk.h" +#include "neug/common/types/property_types.h" #include "neug/main/connection.h" #include "neug/storages/checkpoint_manager.h" #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/graph/edge_table.h" #include "neug/storages/graph/vertex_table.h" +#include "neug/storages/id_indexer.h" #include "neug/storages/loader/loader_utils.h" #include "neug/storages/module/module_broker.h" #include "neug/storages/module/module_factory.h" #include "neug/storages/module/type_name.h" -#include "neug/utils/id_indexer.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" -#include "neug/utils/property/types.h" +#include "neug/storages/property/column.h" +#include "neug/storages/property/table.h" class GeneratedChunkSupplier : public neug::IDataChunkSupplier { public: diff --git a/tests/utils/json_test.cc b/tests/utils/json_test.cc index d526bf189..8fd0d33fc 100644 --- a/tests/utils/json_test.cc +++ b/tests/utils/json_test.cc @@ -134,7 +134,7 @@ TEST_F(JsonTest, TestJsonArray) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); diff --git a/tests/utils/test_reader.cc b/tests/utils/test_reader.cc index 90a4db953..4057283f1 100644 --- a/tests/utils/test_reader.cc +++ b/tests/utils/test_reader.cc @@ -37,7 +37,7 @@ TEST_F(ReaderTest, TestBasicCsvRead) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Verify data: should have 3 columns EXPECT_EQ(ctx.col_num(), 3); @@ -62,7 +62,7 @@ TEST_F(ReaderTest, TestCsvWithTabDelimiter) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -87,7 +87,7 @@ TEST_F(ReaderTest, TestCsvWithCustomQuoting) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -108,7 +108,7 @@ TEST_F(ReaderTest, TestCsvWithNoHeader) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -136,7 +136,7 @@ TEST_F(ReaderTest, TestBatchRead) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Batch mode: data is materialized into Context chunks EXPECT_GT(ctx.chunk_num(), 0); @@ -167,7 +167,7 @@ TEST_F(ReaderTest, TestColumnPruning) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should only have 2 columns (id and score) EXPECT_EQ(ctx.col_num(), 2); @@ -197,7 +197,7 @@ TEST_F(ReaderTest, TestFilterPushdown) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should filter out rows with score <= 90.0 // Expected: Alice (95.5) and Charlie (92.5) - 2 rows @@ -230,7 +230,7 @@ TEST_F(ReaderTest, TestColumnPruningAndFilterPushdown) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should have 2 columns (id, score) and filtered rows (score > 90.0) EXPECT_EQ(ctx.col_num(), 2); @@ -259,7 +259,7 @@ TEST_F(ReaderTest, TestMultipleFiles) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should read all rows from both files (4 rows total) EXPECT_EQ(ctx.col_num(), 3); @@ -285,7 +285,7 @@ TEST_F(ReaderTest, TestForceColumnTypeConversion) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 3); @@ -329,7 +329,7 @@ TEST_F(ReaderTest, TestMultiColumnAndFilterPushdown) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should have 3 columns EXPECT_EQ(ctx.col_num(), 3); @@ -360,7 +360,7 @@ TEST_F(ReaderTest, TestBatchReadWithFilter) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should filter out rows with score <= 90.0 // Expected: Alice (95.5) and Charlie (92.5) - 2 rows @@ -392,7 +392,7 @@ TEST_F(ReaderTest, TestBatchReadWithFilterAndProjection) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should have 2 columns (id, score) and filtered rows (score > 90.0) EXPECT_EQ(ctx.col_num(), 2); @@ -419,7 +419,7 @@ TEST_F(ReaderTest, TestBasicJsonRead) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 3); EXPECT_EQ(ctx.row_num(), 2); @@ -441,7 +441,7 @@ TEST_F(ReaderTest, TestJsonNonExistentColumnThrows) { auto localState = std::make_shared(); execution::Context ctx; - EXPECT_THROW(reader->read(localState, ctx), + EXPECT_THROW(ctx.append_chunks(reader->read(localState)), exception::SchemaMismatchException); } @@ -470,7 +470,7 @@ TEST_F(ReaderTest, TestJsonBatchReadWithFilter) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should filter out rows with score <= 90.0 // Expected: Alice (95.5) and Charlie (92.5) - 2 rows @@ -504,7 +504,7 @@ TEST_F(ReaderTest, TestJsonBatchReadWithFilterAndProjection) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); // Should have 2 columns (id, score) and filtered rows (score > 90.0) EXPECT_EQ(ctx.col_num(), 2); diff --git a/tests/utils/test_table.cc b/tests/utils/test_table.cc index efb3e6ea0..09cc7c214 100644 --- a/tests/utils/test_table.cc +++ b/tests/utils/test_table.cc @@ -21,8 +21,8 @@ #include "neug/storages/checkpoint_manifest.h" #include "neug/storages/module/module_broker.h" #include "neug/storages/module_descriptor.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" +#include "neug/storages/property/column.h" +#include "neug/storages/property/table.h" #include "unittest/utils.h" static const std::vector bool_data = {1, 0, 0, 1, 1, 0, 1, 0, 1, 1}; diff --git a/tests/utils/test_types.cc b/tests/utils/test_types.cc index 4b0b1f4c0..28d5cdc26 100644 --- a/tests/utils/test_types.cc +++ b/tests/utils/test_types.cc @@ -16,8 +16,8 @@ #include #include "neug/common/types/value.h" -#include "neug/utils/property/column.h" -#include "neug/utils/property/table.h" +#include "neug/storages/property/column.h" +#include "neug/storages/property/table.h" #include "neug/utils/serialization/in_archive.h" #include "neug/utils/serialization/out_archive.h" diff --git a/tools/nodejs_bind/src/node_query_request.cc b/tools/nodejs_bind/src/node_query_request.cc index 8bf14a74f..6d2ae8e35 100644 --- a/tools/nodejs_bind/src/node_query_request.cc +++ b/tools/nodejs_bind/src/node_query_request.cc @@ -18,11 +18,11 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/main/query_request.h" #include "neug/utils/access_mode.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" #include diff --git a/tools/nodejs_bind/src/node_query_result.cc b/tools/nodejs_bind/src/node_query_result.cc index 126876c7f..ced13937f 100644 --- a/tools/nodejs_bind/src/node_query_result.cc +++ b/tools/nodejs_bind/src/node_query_result.cc @@ -19,10 +19,10 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/storages/graph/schema.h" #include "neug/utils/bolt_utils.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/tools/python_bind/src/py_query_request.cc b/tools/python_bind/src/py_query_request.cc index 042a38234..fccf36a4f 100644 --- a/tools/python_bind/src/py_query_request.cc +++ b/tools/python_bind/src/py_query_request.cc @@ -17,11 +17,11 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/main/query_request.h" #include "neug/utils/access_mode.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" #include "neug/utils/serialization/in_archive.h" #include diff --git a/tools/python_bind/src/py_query_result.cc b/tools/python_bind/src/py_query_result.cc index d313f9f38..f30a3f64c 100644 --- a/tools/python_bind/src/py_query_result.cc +++ b/tools/python_bind/src/py_query_result.cc @@ -20,11 +20,11 @@ #include #include #include +#include "neug/common/types/property_types.h" #include "neug/compiler/common/arrow/arrow.h" #include "neug/storages/graph/schema.h" #include "neug/utils/bolt_utils.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/types.h" // --------------------------------------------------------------------------- // Arrow C Data Interface helpers – build ArrowSchema / ArrowArray from the From 94a22f96c8e42ce3737bc476b7a4f0c51995094e Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 8 Jul 2026 11:57:01 +0800 Subject: [PATCH 2/8] fix --- .../parquet/src/parquet_export_function.cc | 13 +- include/neug/common/export/export_result.h | 40 ++++ .../common/operators/retrieve/sink.h | 7 - src/compiler/function/csv_export_function.cpp | 17 +- src/compiler/function/json_export_function.cc | 67 +++++-- .../common/operators/retrieve/sink.cc | 35 ++-- src/utils/io/write/writer.cc | 5 + tests/utils/CMakeLists.txt | 3 +- tests/utils/export_result_test.cc | 179 ++++++++++++++++++ 9 files changed, 312 insertions(+), 54 deletions(-) create mode 100644 include/neug/common/export/export_result.h create mode 100644 tests/utils/export_result_test.cc diff --git a/extension/parquet/src/parquet_export_function.cc b/extension/parquet/src/parquet_export_function.cc index 8967b314d..b941d222d 100644 --- a/extension/parquet/src/parquet_export_function.cc +++ b/extension/parquet/src/parquet_export_function.cc @@ -29,12 +29,12 @@ #include "neug/common/columns/list_columns.h" #include "neug/common/columns/struct_columns.h" #include "neug/common/columns/value_columns.h" +#include "neug/common/export/export_result.h" #include "neug/common/types/array_columns.h" #include "neug/common/types/data_chunk.h" #include "neug/common/types/property_types.h" #include "neug/common/types/value.h" #include "neug/compiler/main/metadata_registry.h" -#include "neug/execution/common/operators/retrieve/sink.h" #include "neug/utils/exception/exception.h" #include "parquet/arrow_fs_resolver.h" #include "parquet_options.h" @@ -550,13 +550,10 @@ static execution::Context parquetExecFunc( auto writer = std::make_shared( schema, std::move(arrowFs), entry_schema); - auto source_types = ctx.column_types(); - auto chunks = execution::Sink::materialize_for_export(ctx, graph); - for (const auto& chunk : chunks) { - auto status = writer->write(chunk, source_types); - if (!status.ok()) { - THROW_IO_EXCEPTION("Parquet export failed: " + status.ToString()); - } + auto export_result = neug::materialize_result_for_export(ctx, graph); + auto status = writer->write(export_result.chunk, export_result.source_types); + if (!status.ok()) { + THROW_IO_EXCEPTION("Parquet export failed: " + status.ToString()); } LOG(INFO) << "[Parquet Export] Export completed successfully"; ctx.clear(); diff --git a/include/neug/common/export/export_result.h b/include/neug/common/export/export_result.h new file mode 100644 index 000000000..664ce7f38 --- /dev/null +++ b/include/neug/common/export/export_result.h @@ -0,0 +1,40 @@ +/** 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 "neug/common/types.h" +#include "neug/common/types/data_chunk.h" + +namespace neug { + +class StorageReadInterface; + +namespace execution { +class Context; +} + +struct ExportResult { + DataChunk chunk; + std::vector source_types; +}; + +/// Materializes graph values for file export and merges Context chunks into a +/// single DataChunk so writers can emit one complete output artifact. +ExportResult materialize_result_for_export(const execution::Context& ctx, + const StorageReadInterface& graph); + +} // namespace neug diff --git a/include/neug/execution/common/operators/retrieve/sink.h b/include/neug/execution/common/operators/retrieve/sink.h index 211701d64..eabfa59ff 100644 --- a/include/neug/execution/common/operators/retrieve/sink.h +++ b/include/neug/execution/common/operators/retrieve/sink.h @@ -14,8 +14,6 @@ */ #pragma once -#include "neug/common/types.h" -#include "neug/common/types/data_chunk.h" #include "neug/generated/proto/response/response.pb.h" #include "neug/main/query_result.h" namespace neug { @@ -29,11 +27,6 @@ class Context; class Sink { public: - /// Materialize graph-type columns to VARCHAR strings per Context chunk. - /// Returns one DataChunk per Context chunk without cross-chunk column fusion. - static std::vector materialize_for_export( - const Context& ctx, const StorageReadInterface& graph); - static void sink_results(const Context& ctx, const StorageReadInterface& graph, neug::QueryResponse* response); diff --git a/src/compiler/function/csv_export_function.cpp b/src/compiler/function/csv_export_function.cpp index c9b25ba37..8938bfecb 100644 --- a/src/compiler/function/csv_export_function.cpp +++ b/src/compiler/function/csv_export_function.cpp @@ -20,9 +20,9 @@ * Zhou Xiaoli in 2025 to support Neug-specific features. */ +#include "neug/common/export/export_result.h" #include "neug/compiler/function/export/export_function.h" #include "neug/compiler/main/metadata_registry.h" -#include "neug/execution/common/operators/retrieve/sink.h" #include "neug/utils/io/write/writer.h" namespace neug { @@ -63,16 +63,13 @@ execution::Context writeExecFunc( convertFileSchemaOptions(schema); auto writer = std::make_shared( schema, entry_schema); - auto source_types = ctx.column_types(); - auto chunks = neug::execution::Sink::materialize_for_export(ctx, graph); - for (const auto& chunk : chunks) { - auto status = writer->write(chunk, source_types); - if (!status.ok()) { - if (status.error_code() == StatusCode::ERR_PERMISSION) { - THROW_PERMISSION_DENIED("Export failed: " + status.ToString()); - } - THROW_IO_EXCEPTION("Export failed: " + status.ToString()); + auto export_result = neug::materialize_result_for_export(ctx, graph); + auto status = writer->write(export_result.chunk, export_result.source_types); + if (!status.ok()) { + if (status.error_code() == StatusCode::ERR_PERMISSION) { + THROW_PERMISSION_DENIED("Export failed: " + status.ToString()); } + THROW_IO_EXCEPTION("Export failed: " + status.ToString()); } ctx.clear(); return ctx; diff --git a/src/compiler/function/json_export_function.cc b/src/compiler/function/json_export_function.cc index a220ca78a..7f4de39bc 100644 --- a/src/compiler/function/json_export_function.cc +++ b/src/compiler/function/json_export_function.cc @@ -19,9 +19,9 @@ #include #include #include +#include "neug/common/export/export_result.h" #include "neug/common/types/data_chunk.h" #include "neug/common/types/value.h" -#include "neug/execution/common/operators/retrieve/sink.h" #include "neug/utils/io/stream/output_stream.h" #include @@ -108,6 +108,49 @@ static neug::result valueToJsonValue( allocator); return v; } + case DataTypeId::kList: { + rapidjson::Value arr(rapidjson::kArrayType); + const auto& children = ListValue::GetChildren(value); + for (const auto& child : children) { + auto child_json = valueToJsonValue(child, doc); + if (!child_json) { + return tl::make_unexpected(child_json.error()); + } + arr.PushBack(std::move(*child_json), allocator); + } + return arr; + } + case DataTypeId::kArray: { + rapidjson::Value arr(rapidjson::kArrayType); + const auto& children = ArrayValue::GetChildren(value); + for (const auto& child : children) { + auto child_json = valueToJsonValue(child, doc); + if (!child_json) { + return tl::make_unexpected(child_json.error()); + } + arr.PushBack(std::move(*child_json), allocator); + } + return arr; + } + case DataTypeId::kStruct: { + rapidjson::Value obj(rapidjson::kObjectType); + const auto& children = StructValue::GetChildren(value); + const auto& field_names = StructType::GetFieldNames(value.type()); + for (size_t i = 0; i < children.size(); ++i) { + const auto field_name = i < field_names.size() + ? field_names[i] + : ("field_" + std::to_string(i)); + rapidjson::Value key(field_name.c_str(), + static_cast(field_name.size()), + allocator); + auto child_json = valueToJsonValue(children[i], doc); + if (!child_json) { + return tl::make_unexpected(child_json.error()); + } + obj.AddMember(key, std::move(*child_json), allocator); + } + return obj; + } default: { const auto& s = value.to_string(); rapidjson::Value v; @@ -304,13 +347,10 @@ static execution::Context jsonExecFunc( } auto writer = std::make_shared( schema, entry_schema); - auto source_types = ctx.column_types(); - auto chunks = neug::execution::Sink::materialize_for_export(ctx, graph); - for (const auto& chunk : chunks) { - auto status = writer->write(chunk, source_types); - if (!status.ok()) { - THROW_IO_EXCEPTION("Export failed: " + status.ToString()); - } + auto export_result = neug::materialize_result_for_export(ctx, graph); + auto status = writer->write(export_result.chunk, export_result.source_types); + if (!status.ok()) { + THROW_IO_EXCEPTION("Export failed: " + status.ToString()); } ctx.clear(); return ctx; @@ -343,13 +383,10 @@ static execution::Context jsonLExecFunc( } auto writer = std::make_shared(schema, entry_schema); - auto source_types = ctx.column_types(); - auto chunks = neug::execution::Sink::materialize_for_export(ctx, graph); - for (const auto& chunk : chunks) { - auto status = writer->write(chunk, source_types); - if (!status.ok()) { - THROW_IO_EXCEPTION("Export failed: " + status.ToString()); - } + auto export_result = neug::materialize_result_for_export(ctx, graph); + auto status = writer->write(export_result.chunk, export_result.source_types); + if (!status.ok()) { + THROW_IO_EXCEPTION("Export failed: " + status.ToString()); } ctx.clear(); return ctx; diff --git a/src/execution/common/operators/retrieve/sink.cc b/src/execution/common/operators/retrieve/sink.cc index 2e02c8fc6..b0d0432d7 100644 --- a/src/execution/common/operators/retrieve/sink.cc +++ b/src/execution/common/operators/retrieve/sink.cc @@ -22,6 +22,7 @@ #include "neug/common/columns/struct_columns.h" #include "neug/common/columns/value_columns.h" #include "neug/common/columns/vertex_columns.h" +#include "neug/common/export/export_result.h" #include "neug/common/types/array_columns.h" #include "neug/common/types/data_chunk.h" #include "neug/common/types/property_types.h" @@ -692,25 +693,33 @@ std::shared_ptr materialize_column_for_export( } // namespace -std::vector Sink::materialize_for_export( - const Context& ctx, const StorageReadInterface& graph) { - std::vector result; - result.reserve(ctx.chunk_num()); - for (size_t c = 0; c < ctx.chunk_num(); ++c) { - const auto& ctx_chunk = ctx.chunk(c); - DataChunk chunk; - int alias = 0; - for (size_t tag_id : ctx.tag_ids) { - auto col = ctx_chunk.get(tag_id); +} // namespace execution + +ExportResult materialize_result_for_export(const execution::Context& ctx, + const StorageReadInterface& graph) { + ExportResult result; + int alias = 0; + for (size_t tag_id : ctx.tag_ids) { + std::shared_ptr merged; + for (size_t c = 0; c < ctx.chunk_num(); ++c) { + auto col = ctx.chunk(c).get(tag_id); if (col == nullptr) { continue; } - chunk.set(alias++, materialize_column_for_export(col, graph)); + if (merged == nullptr) { + merged = col; + } else { + merged = merged->union_col(col); + } + } + if (merged == nullptr) { + continue; } - result.push_back(std::move(chunk)); + result.source_types.push_back(merged->elem_type()); + result.chunk.set(alias++, + execution::materialize_column_for_export(merged, graph)); } return result; } -} // namespace execution } // namespace neug diff --git a/src/utils/io/write/writer.cc b/src/utils/io/write/writer.cc index 8c3680698..aacd4ca11 100644 --- a/src/utils/io/write/writer.cc +++ b/src/utils/io/write/writer.cc @@ -134,6 +134,11 @@ neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( write(reinterpret_cast("e_char_), sizeof(char)); return neug::Status::OK(); } + if (type_id == DataTypeId::kBoolean) { + const auto* str = value.GetValue() ? "true" : "false"; + write(reinterpret_cast(str), strlen(str)); + return neug::Status::OK(); + } const auto& str = value.to_string(); write(reinterpret_cast(str.c_str()), str.size()); return neug::Status::OK(); diff --git a/tests/utils/CMakeLists.txt b/tests/utils/CMakeLists.txt index 976a54c64..6a4847fcc 100644 --- a/tests/utils/CMakeLists.txt +++ b/tests/utils/CMakeLists.txt @@ -7,4 +7,5 @@ add_neug_test( test_exception.cc test_sniffer.cc test_file_utils.cc - json_test.cc) + json_test.cc + export_result_test.cc) diff --git a/tests/utils/export_result_test.cc b/tests/utils/export_result_test.cc new file mode 100644 index 000000000..9b9430ef8 --- /dev/null +++ b/tests/utils/export_result_test.cc @@ -0,0 +1,179 @@ +/** + * 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 +#include +#include +#include +#include +#include + +#include "neug/common/columns/struct_columns.h" +#include "neug/common/columns/value_columns.h" +#include "neug/common/export/export_result.h" +#include "neug/common/types/data_chunk.h" +#include "neug/common/types/value.h" +#include "neug/compiler/function/export/json_export_function.h" +#include "neug/execution/common/context.h" +#include "neug/storages/graph/graph_interface.h" +#include "neug/storages/graph/graph_view.h" +#include "neug/storages/graph/property_graph.h" +#include "neug/utils/io/write/writer.h" + +namespace neug { +namespace test { +namespace { + +constexpr const char* EXPORT_RESULT_TEST_DIR = "/tmp/neug_export_result_test"; + +std::shared_ptr intColumn(std::vector values) { + ValueColumnBuilder builder; + builder.reserve(values.size()); + for (auto value : values) { + builder.push_back_opt(value); + } + return builder.finish(); +} + +std::shared_ptr stringColumn( + const std::vector& values) { + ValueColumnBuilder builder; + builder.reserve(values.size()); + for (const auto& value : values) { + builder.push_back_opt(value); + } + return builder.finish(); +} + +std::shared_ptr boolColumn(std::vector values) { + ValueColumnBuilder builder; + builder.reserve(values.size()); + for (auto value : values) { + builder.push_back_opt(value); + } + return builder.finish(); +} + +std::vector readLines(const std::string& path) { + std::ifstream file(path); + std::vector lines; + std::string line; + while (std::getline(file, line)) { + lines.push_back(line); + } + return lines; +} + +class ExportResultTest : public ::testing::Test { + public: + void SetUp() override { + std::filesystem::remove_all(EXPORT_RESULT_TEST_DIR); + std::filesystem::create_directories(EXPORT_RESULT_TEST_DIR); + } + + void TearDown() override { + std::filesystem::remove_all(EXPORT_RESULT_TEST_DIR); + } +}; + +TEST_F(ExportResultTest, MaterializerMergesChunksBeforeCsvWrite) { + execution::Context ctx; + ctx.tag_ids = {0, 1, 2}; + + DataChunk first; + first.set(0, intColumn({1, 2})); + first.set(1, stringColumn({"alice", "bob"})); + first.set(2, boolColumn({true, false})); + ctx.append_chunk(std::move(first)); + + DataChunk second; + second.set(0, intColumn({3})); + second.set(1, stringColumn({"carol"})); + second.set(2, boolColumn({true})); + ctx.append_chunk(std::move(second)); + + PropertyGraph graph; + GraphView view(graph); + StorageReadInterface reader(view, 1); + + auto export_result = materialize_result_for_export(ctx, reader); + ASSERT_EQ(export_result.chunk.col_num(), 3); + ASSERT_EQ(export_result.chunk.row_num(), 3); + ASSERT_EQ(export_result.source_types.size(), 3); + EXPECT_EQ(export_result.source_types[0].id(), DataTypeId::kInt32); + EXPECT_EQ(export_result.source_types[1].id(), DataTypeId::kVarchar); + EXPECT_EQ(export_result.source_types[2].id(), DataTypeId::kBoolean); + + reader::FileSchema schema; + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/merged.csv"}; + schema.format = "csv"; + auto entry_schema = std::make_shared(); + entry_schema->columnNames = {"id", "name", "active"}; + + writer::CsvQueryExportWriter writer(schema, entry_schema); + auto status = writer.write(export_result.chunk, export_result.source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + + auto lines = readLines(schema.paths[0]); + ASSERT_EQ(lines.size(), 4); + EXPECT_EQ(lines[0], "id|name|active"); + EXPECT_EQ(lines[1], "1|\"alice\"|true"); + EXPECT_EQ(lines[2], "2|\"bob\"|false"); + EXPECT_EQ(lines[3], "3|\"carol\"|true"); +} + +TEST_F(ExportResultTest, JsonArrayWriterEmitsNestedValues) { + auto list_type = DataType::List(DataType(DataTypeId::kInt32)); + std::vector field_names = {"a", "items"}; + std::vector field_types = {DataType(DataTypeId::kInt32), list_type}; + auto struct_type = + DataType::Struct(std::move(field_names), std::move(field_types)); + + StructColumnBuilder builder(struct_type); + std::vector list_values; + list_values.push_back(Value::INT32(1)); + list_values.push_back(Value::INT32(2)); + + std::vector children; + children.push_back(Value::INT32(7)); + children.push_back( + Value::LIST(DataType(DataTypeId::kInt32), std::move(list_values))); + builder.push_back_elem(Value::STRUCT(struct_type, std::move(children))); + + DataChunk chunk; + chunk.set(0, builder.finish()); + + reader::FileSchema schema; + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/nested.json"}; + schema.format = "json"; + auto entry_schema = std::make_shared(); + entry_schema->columnNames = {"payload"}; + + writer::JsonArrayExportWriter writer(schema, entry_schema); + auto status = writer.write(chunk, {struct_type}); + ASSERT_TRUE(status.ok()) << status.ToString(); + + std::ifstream file(schema.paths[0]); + std::string json((std::istreambuf_iterator(file)), + std::istreambuf_iterator()); + EXPECT_EQ(json, R"([{"payload":{"a":7,"items":[1,2]}}])"); +} + +} // namespace +} // namespace test +} // namespace neug From fd1e9f0d69f0fe4b554fe6641610fd9163a8afa2 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Wed, 8 Jul 2026 14:42:42 +0800 Subject: [PATCH 3/8] fix --- include/neug/utils/io/write/writer.h | 8 +++++-- src/utils/io/write/writer.cc | 32 +++++++++++++++++++++++----- 2 files changed, 33 insertions(+), 7 deletions(-) diff --git a/include/neug/utils/io/write/writer.h b/include/neug/utils/io/write/writer.h index 467dd160f..cb560e054 100644 --- a/include/neug/utils/io/write/writer.h +++ b/include/neug/utils/io/write/writer.h @@ -70,7 +70,8 @@ class DataChunkCSVStringFormatBuffer { public: DataChunkCSVStringFormatBuffer(const DataChunk& chunk, const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema); + const reader::EntrySchema& entry_schema, + const std::vector& source_types); ~DataChunkCSVStringFormatBuffer() = default; void addValue(size_t row_idx, size_t col_idx); @@ -78,13 +79,16 @@ class DataChunkCSVStringFormatBuffer { neug::Status flush(io::OutputStream& stream); private: - neug::Status formatValueToStr(const Value& value, size_t row_idx); + neug::Status formatValueToStr(const Value& value, size_t row_idx, + size_t col_idx); + bool shouldWriteRawString(size_t col_idx) const; void writeWithEscapes(char* toEscape, char escape, const std::string& str); void write(const uint8_t* buffer, uint64_t len); const DataChunk& chunk_; const reader::FileSchema& schema_; const reader::EntrySchema& entry_schema_; + const std::vector& source_types_; BinaryData blob_; size_t capacity_; uint8_t* data_; diff --git a/src/utils/io/write/writer.cc b/src/utils/io/write/writer.cc index aacd4ca11..9a2dd6874 100644 --- a/src/utils/io/write/writer.cc +++ b/src/utils/io/write/writer.cc @@ -32,10 +32,12 @@ namespace writer { DataChunkCSVStringFormatBuffer::DataChunkCSVStringFormatBuffer( const DataChunk& chunk, const reader::FileSchema& schema, - const reader::EntrySchema& entry_schema) + const reader::EntrySchema& entry_schema, + const std::vector& source_types) : chunk_(chunk), schema_(schema), entry_schema_(entry_schema), + source_types_(source_types), capacity_(DEFAULT_CAPACITY) { WriteOptions write_opts; size_t batch_size = write_opts.batch_rows.get(schema.options); @@ -120,7 +122,7 @@ void DataChunkCSVStringFormatBuffer::write(const uint8_t* buffer, } neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( - const Value& value, size_t row_idx) { + const Value& value, size_t row_idx, size_t col_idx) { if (value.IsNull()) { return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, "Value is invalid, rowIdx=" + std::to_string(row_idx)); @@ -128,6 +130,10 @@ neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( const auto type_id = value.type().id(); if (type_id == DataTypeId::kVarchar) { const auto& str = StringValue::Get(value); + if (shouldWriteRawString(col_idx)) { + write(reinterpret_cast(str.c_str()), str.size()); + return neug::Status::OK(); + } write(reinterpret_cast("e_char_), sizeof(char)); char escape_chars[] = {escape_char_, quote_char_}; writeWithEscapes(escape_chars, escape_char_, str); @@ -144,6 +150,20 @@ neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( return neug::Status::OK(); } +bool DataChunkCSVStringFormatBuffer::shouldWriteRawString(size_t col_idx) const { + if (col_idx >= source_types_.size()) { + return false; + } + switch (source_types_[col_idx].id()) { + case DataTypeId::kVertex: + case DataTypeId::kEdge: + case DataTypeId::kPath: + return true; + default: + return false; + } +} + void DataChunkCSVStringFormatBuffer::addValue(size_t row_idx, size_t col_idx) { if (col_idx >= chunk_.col_num() || chunk_.columns[col_idx] == nullptr) { THROW_IO_EXCEPTION("Column index out of range: colIdx=" + @@ -165,7 +185,8 @@ void DataChunkCSVStringFormatBuffer::addValue(size_t row_idx, size_t col_idx) { write(reinterpret_cast(DEFAULT_NULL_STR), strlen(DEFAULT_NULL_STR)); } else { - auto str_result = formatValueToStr(column->get_elem(row_idx), row_idx); + auto str_result = formatValueToStr(column->get_elem(row_idx), row_idx, + col_idx); if (!str_result.ok()) { if (!ignore_errors_) { THROW_IO_EXCEPTION( @@ -194,7 +215,7 @@ neug::Status DataChunkCSVStringFormatBuffer::flush(io::OutputStream& stream) { } neug::Status CsvQueryExportWriter::write( - const DataChunk& chunk, const std::vector& /*source_types*/) { + const DataChunk& chunk, const std::vector& source_types) { if (schema_.paths.empty()) { return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, "Schema paths is empty"); @@ -215,7 +236,8 @@ neug::Status CsvQueryExportWriter::write( "Batch size should be positive"); } auto csv_buffer = - DataChunkCSVStringFormatBuffer(chunk, schema_, *entry_schema_); + DataChunkCSVStringFormatBuffer(chunk, schema_, *entry_schema_, + source_types); csv_buffer.addHeader(); for (size_t i = 0; i < chunk.row_num(); ++i) { for (size_t j = 0; j < chunk.col_num(); ++j) { From 01db079f4ccc9ef3771e6fb96d3959002a64815a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=99=93=E7=AE=80?= Date: Wed, 8 Jul 2026 15:47:10 +0800 Subject: [PATCH 4/8] fmt --- src/utils/io/write/writer.cc | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/utils/io/write/writer.cc b/src/utils/io/write/writer.cc index 9a2dd6874..a95700e60 100644 --- a/src/utils/io/write/writer.cc +++ b/src/utils/io/write/writer.cc @@ -150,7 +150,8 @@ neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( return neug::Status::OK(); } -bool DataChunkCSVStringFormatBuffer::shouldWriteRawString(size_t col_idx) const { +bool DataChunkCSVStringFormatBuffer::shouldWriteRawString( + size_t col_idx) const { if (col_idx >= source_types_.size()) { return false; } @@ -185,8 +186,8 @@ void DataChunkCSVStringFormatBuffer::addValue(size_t row_idx, size_t col_idx) { write(reinterpret_cast(DEFAULT_NULL_STR), strlen(DEFAULT_NULL_STR)); } else { - auto str_result = formatValueToStr(column->get_elem(row_idx), row_idx, - col_idx); + auto str_result = + formatValueToStr(column->get_elem(row_idx), row_idx, col_idx); if (!str_result.ok()) { if (!ignore_errors_) { THROW_IO_EXCEPTION( @@ -235,9 +236,8 @@ neug::Status CsvQueryExportWriter::write( return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, "Batch size should be positive"); } - auto csv_buffer = - DataChunkCSVStringFormatBuffer(chunk, schema_, *entry_schema_, - source_types); + auto csv_buffer = DataChunkCSVStringFormatBuffer( + chunk, schema_, *entry_schema_, source_types); csv_buffer.addHeader(); for (size_t i = 0; i < chunk.row_num(); ++i) { for (size_t j = 0; j < chunk.col_num(); ++j) { From 357171571aec40703b574843e9f9b7327cda50e3 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 24 Jul 2026 17:25:42 +0800 Subject: [PATCH 5/8] fix: integrate module refactor with latest main --- extension/parquet/src/arrow_column.cc | 2 +- extension/parquet/src/arrow_options.cc | 5 +- extension/parquet/src/arrow_reader.cc | 55 +++++++++++++++++-- include/neug/common/types/property_types.h | 2 +- .../execution/execute/ops/ddl/ddl_utils.h | 2 +- include/neug/storages/graph/dirty_tracker.h | 2 +- .../neug/utils/property/property_definition.h | 2 +- src/compiler/binder/bind/bind_ddl.cpp | 2 +- .../binder/bind/copy/bind_copy_from.cpp | 2 +- tests/storage/test_checkpoint.cc | 2 +- tests/utils/json_test.cc | 6 +- tools/python_bind/tests/test_load_array.py | 23 ++++---- 12 files changed, 79 insertions(+), 26 deletions(-) diff --git a/extension/parquet/src/arrow_column.cc b/extension/parquet/src/arrow_column.cc index 4d5559cd2..ebeb83006 100644 --- a/extension/parquet/src/arrow_column.cc +++ b/extension/parquet/src/arrow_column.cc @@ -146,7 +146,7 @@ static Value arrow_value_at(const arrow::Array& array, int64_t index, Interval interval; interval.from_mill_seconds( arrow_time_to_milliseconds(value, duration_type.unit())); - return Value::INTERVAL(interval); + return Value::INTERVAL(interval.normalize()); } case arrow::Type::FIXED_SIZE_LIST: { const auto& list = static_cast(array); diff --git a/extension/parquet/src/arrow_options.cc b/extension/parquet/src/arrow_options.cc index a87e14ea3..d0d1ed5da 100644 --- a/extension/parquet/src/arrow_options.cc +++ b/extension/parquet/src/arrow_options.cc @@ -69,7 +69,10 @@ bool ArrowOptionsBuilder::projectColumns(ArrowOptions& options) { } } - auto dataset_schema = createSchema(entrySchema); + auto dataset_schema = options.scanOptions->dataset_schema; + if (!dataset_schema) { + dataset_schema = createSchema(entrySchema); + } auto project_desc = arrow::dataset::ProjectionDescr::FromNames(columns, *dataset_schema); if (!project_desc.ok()) { diff --git a/extension/parquet/src/arrow_reader.cc b/extension/parquet/src/arrow_reader.cc index 3b4c18752..9c25b0821 100644 --- a/extension/parquet/src/arrow_reader.cc +++ b/extension/parquet/src/arrow_reader.cc @@ -31,6 +31,51 @@ namespace neug { namespace reader { +static std::shared_ptr reconcileIntervalType( + const std::shared_ptr& expected, + const std::shared_ptr& file_type) { + if ((expected->id() == arrow::Type::STRING || + expected->id() == arrow::Type::LARGE_STRING) && + file_type->id() == arrow::Type::DURATION) { + return file_type; + } + if (expected->id() != arrow::Type::FIXED_SIZE_LIST || + file_type->id() != arrow::Type::FIXED_SIZE_LIST) { + return expected; + } + + const auto& expected_list = + static_cast(*expected); + const auto& file_list = + static_cast(*file_type); + if (expected_list.list_size() != file_list.list_size()) { + return expected; + } + auto child = + reconcileIntervalType(expected_list.value_type(), file_list.value_type()); + if (child->Equals(expected_list.value_type())) { + return expected; + } + return arrow::fixed_size_list(child, expected_list.list_size()); +} + +static std::shared_ptr reconcileIntervalSchema( + const std::shared_ptr& expected, + const std::shared_ptr& file_schema) { + std::vector> fields; + fields.reserve(expected->num_fields()); + for (const auto& field : expected->fields()) { + auto file_field = file_schema->GetFieldByName(field->name()); + if (!file_field) { + fields.push_back(field); + continue; + } + auto type = reconcileIntervalType(field->type(), file_field->type()); + fields.push_back(field->WithType(std::move(type))); + } + return arrow::schema(std::move(fields), expected->metadata()); +} + std::vector ArrowReader::read( std::shared_ptr localState) { if (!sharedState) { @@ -79,10 +124,6 @@ std::shared_ptr ArrowReader::createScanner( THROW_INVALID_ARGUMENT_EXCEPTION("Failed to build arrow options"); } - if (!optionsBuilder->projectColumns(arrowOptions)) { - LOG(WARNING) << "Failed to set column projection, using all columns"; - } - if (!optionsBuilder->skipRows(arrowOptions)) { LOG(WARNING) << "Failed to set row filter, using no filter"; } @@ -108,6 +149,8 @@ std::shared_ptr ArrowReader::createScanner( fileSchema->ToString()); } } + scan_opts->dataset_schema = + reconcileIntervalSchema(scan_opts->dataset_schema, fileSchema); } dataset_result = factory->Finish(scan_opts->dataset_schema); } else { @@ -123,6 +166,10 @@ std::shared_ptr ArrowReader::createScanner( } auto dataset = dataset_result.ValueOrDie(); + if (!optionsBuilder->projectColumns(arrowOptions)) { + LOG(WARNING) << "Failed to set column projection, using all columns"; + } + arrow::dataset::ScannerBuilder scanner_builder(dataset, scan_opts); auto scanner_result = scanner_builder.Finish(); if (!scanner_result.ok()) { diff --git a/include/neug/common/types/property_types.h b/include/neug/common/types/property_types.h index 5cc6dd6f3..d60b8b3f7 100644 --- a/include/neug/common/types/property_types.h +++ b/include/neug/common/types/property_types.h @@ -187,7 +187,7 @@ struct Interval { // Normalize to interval bounds. inline static void borrow(const int64_t msf, int64_t& lsf, int32_t& f, const int64_t scale); - inline Interval normalize() const; + Interval normalize() const; inline Interval& operator=(const Interval& rhs) = default; diff --git a/include/neug/execution/execute/ops/ddl/ddl_utils.h b/include/neug/execution/execute/ops/ddl/ddl_utils.h index 32b8eb236..6857aa75a 100644 --- a/include/neug/execution/execute/ops/ddl/ddl_utils.h +++ b/include/neug/execution/execute/ops/ddl/ddl_utils.h @@ -16,8 +16,8 @@ #include +#include "neug/common/types/property_types.h" #include "neug/storages/graph/schema.h" -#include "neug/utils/property/types.h" #include "neug/utils/result.h" namespace neug { diff --git a/include/neug/storages/graph/dirty_tracker.h b/include/neug/storages/graph/dirty_tracker.h index 3d21b97ff..b3479ce05 100644 --- a/include/neug/storages/graph/dirty_tracker.h +++ b/include/neug/storages/graph/dirty_tracker.h @@ -20,7 +20,7 @@ #include #include -#include "neug/utils/property/types.h" +#include "neug/common/types/property_types.h" namespace neug { diff --git a/include/neug/utils/property/property_definition.h b/include/neug/utils/property/property_definition.h index 20672bd49..275760540 100644 --- a/include/neug/utils/property/property_definition.h +++ b/include/neug/utils/property/property_definition.h @@ -18,8 +18,8 @@ #include #include +#include "neug/common/types/property_types.h" #include "neug/common/types/value.h" -#include "neug/utils/property/types.h" namespace neug { diff --git a/src/compiler/binder/bind/bind_ddl.cpp b/src/compiler/binder/bind/bind_ddl.cpp index e6506adcf..8cab407ad 100644 --- a/src/compiler/binder/bind/bind_ddl.cpp +++ b/src/compiler/binder/bind/bind_ddl.cpp @@ -23,6 +23,7 @@ #include "neug/compiler/binder/binder.h" #include +#include "neug/common/types/default_value.h" #include "neug/compiler/binder/ddl/bound_alter.h" #include "neug/compiler/binder/ddl/bound_create_sequence.h" #include "neug/compiler/binder/ddl/bound_create_table.h" @@ -51,7 +52,6 @@ #include "neug/compiler/parser/expression/parsed_literal_expression.h" #include "neug/utils/exception/exception.h" #include "neug/utils/exception/message.h" -#include "neug/utils/property/default_value.h" using namespace neug::common; using namespace neug::parser; diff --git a/src/compiler/binder/bind/copy/bind_copy_from.cpp b/src/compiler/binder/bind/copy/bind_copy_from.cpp index 6f5916ab5..7f5819b6d 100644 --- a/src/compiler/binder/bind/copy/bind_copy_from.cpp +++ b/src/compiler/binder/bind/copy/bind_copy_from.cpp @@ -23,6 +23,7 @@ #include "neug/compiler/binder/binder.h" #include +#include "neug/common/types/default_value.h" #include "neug/compiler/binder/copy/bound_copy_from.h" #include "neug/compiler/binder/ddl/bound_create_table_info.h" #include "neug/compiler/binder/expression_binder.h" @@ -45,7 +46,6 @@ #include "neug/compiler/parser/expression/parsed_literal_expression.h" #include "neug/compiler/parser/scan_source.h" #include "neug/utils/exception/exception.h" -#include "neug/utils/property/default_value.h" #include "neug/utils/property/property_definition.h" using namespace neug::binder; diff --git a/tests/storage/test_checkpoint.cc b/tests/storage/test_checkpoint.cc index 57ee8e538..5078fcb39 100644 --- a/tests/storage/test_checkpoint.cc +++ b/tests/storage/test_checkpoint.cc @@ -45,7 +45,7 @@ #include "neug/storages/graph/property_graph.h" #include "neug/storages/graph/schema.h" #include "neug/storages/module_descriptor.h" -#include "neug/utils/property/column.h" +#include "neug/storages/property/column.h" #include "unittest/utils.h" namespace { diff --git a/tests/utils/json_test.cc b/tests/utils/json_test.cc index 8fd0d33fc..ae4577198 100644 --- a/tests/utils/json_test.cc +++ b/tests/utils/json_test.cc @@ -161,7 +161,7 @@ TEST_F(JsonTest, TestJsonArrayColumn) { auto localState = std::make_shared(); execution::Context ctx; - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); EXPECT_EQ(ctx.col_num(), 2); EXPECT_EQ(ctx.row_num(), 2); @@ -187,7 +187,7 @@ TEST_F(JsonTest, TestJsonArrayColumnLengthMismatch) { execution::Context ctx; try { - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); FAIL() << "Expected an ARRAY length mismatch"; } catch (const std::exception& error) { EXPECT_NE( @@ -208,7 +208,7 @@ TEST_F(JsonTest, TestJsonArrayColumnRejectsNonArray) { execution::Context ctx; try { - reader->read(localState, ctx); + ctx.append_chunks(reader->read(localState)); FAIL() << "Expected a non-array conversion error"; } catch (const std::exception& error) { EXPECT_NE(std::string(error.what()) diff --git a/tools/python_bind/tests/test_load_array.py b/tools/python_bind/tests/test_load_array.py index 9142f6b04..09c60c7a3 100644 --- a/tools/python_bind/tests/test_load_array.py +++ b/tools/python_bind/tests/test_load_array.py @@ -178,16 +178,19 @@ def test_parquet_nested_string_array(self): def test_parquet_null_array(self): """LOAD FROM Parquet preserves a null fixed-size array.""" pa = pytest.importorskip("pyarrow") - parquet_path = self._write_parquet( - "null_array.parquet", - { - "id": pa.array([1, 2, 3], type=pa.int64()), - "values": pa.array( - [[1.0, 2.0, 3.0], None, [7.0, 8.0, 9.0]], - type=pa.list_(pa.float32(), 3), - ), - }, - ) + try: + parquet_path = self._write_parquet( + "null_array.parquet", + { + "id": pa.array([1, 2, 3], type=pa.int64()), + "values": pa.array( + [[1.0, 2.0, 3.0], None, [7.0, 8.0, 9.0]], + type=pa.list_(pa.float32(), 3), + ), + }, + ) + except pa.ArrowNotImplementedError: + pytest.skip("PyArrow cannot write a null fixed-size list to Parquet") self.conn.execute("LOAD PARQUET") result = list( self.conn.execute( From 3771acbfc2229be6f4c41a4e9ada0fe4f739a89e Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 24 Jul 2026 18:46:37 +0800 Subject: [PATCH 6/8] fix: preserve nested graph values during export --- include/neug/common/columns/list_columns.h | 6 +- include/neug/common/types/array_columns.h | 6 +- src/compiler/function/json_export_function.cc | 69 ++++++++++------ .../common/operators/retrieve/sink.cc | 72 ++++++++++++++--- tests/utils/export_result_test.cc | 81 +++++++++++++++++++ tools/python_bind/tests/test_export.py | 30 +++++++ 6 files changed, 223 insertions(+), 41 deletions(-) diff --git a/include/neug/common/columns/list_columns.h b/include/neug/common/columns/list_columns.h index 0f62e3793..31b7da9ec 100644 --- a/include/neug/common/columns/list_columns.h +++ b/include/neug/common/columns/list_columns.h @@ -122,7 +122,11 @@ class ListColumnBuilder : public IContextColumnBuilder { assert(val.type().id() == DataTypeId::kList); const auto& values = ListValue::GetChildren(val); for (const auto& v : values) { - child_builder_->push_back_elem(v); + if (v.IsNull()) { + child_builder_->push_back_null(); + } else { + child_builder_->push_back_elem(v); + } } list_item item = {cur_offset_, values.size()}; items_.push_back(item); diff --git a/include/neug/common/types/array_columns.h b/include/neug/common/types/array_columns.h index 1e548e277..9b6a7496a 100644 --- a/include/neug/common/types/array_columns.h +++ b/include/neug/common/types/array_columns.h @@ -118,7 +118,11 @@ class ContextArrayColumnBuilder : public IContextColumnBuilder { " elements, got " + std::to_string(children.size())); } for (const auto& v : children) { - child_builder_->push_back_elem(v); + if (v.IsNull()) { + child_builder_->push_back_null(); + } else { + child_builder_->push_back_elem(v); + } } } diff --git a/src/compiler/function/json_export_function.cc b/src/compiler/function/json_export_function.cc index 7f4de39bc..e8770163e 100644 --- a/src/compiler/function/json_export_function.cc +++ b/src/compiler/function/json_export_function.cc @@ -59,12 +59,26 @@ static neug::result parseJsonStringToValue( return v; } +static bool isSerializedGraphJson(DataTypeId id) { + return id == DataTypeId::kVertex || id == DataTypeId::kEdge || + id == DataTypeId::kPath; +} + static neug::result valueToJsonValue( - const Value& value, rapidjson::Document& doc) { + const Value& value, const DataType* source_type, int row, + rapidjson::Document& doc) { auto& allocator = doc.GetAllocator(); if (value.IsNull()) { return rapidjson::Value(rapidjson::kNullType); } + if (source_type != nullptr && isSerializedGraphJson(source_type->id())) { + const auto& str = StringValue::Get(value); + if (str.empty()) { + return rapidjson::Value(rapidjson::kNullType); + } + return parseJsonStringToValue(str, row, doc, + source_type->ToString().c_str()); + } switch (value.type().id()) { case DataTypeId::kBoolean: return rapidjson::Value(value.GetValue()); @@ -111,8 +125,12 @@ static neug::result valueToJsonValue( case DataTypeId::kList: { rapidjson::Value arr(rapidjson::kArrayType); const auto& children = ListValue::GetChildren(value); + const DataType* source_child_type = nullptr; + if (source_type != nullptr && source_type->id() == DataTypeId::kList) { + source_child_type = &ListType::GetChildType(*source_type); + } for (const auto& child : children) { - auto child_json = valueToJsonValue(child, doc); + auto child_json = valueToJsonValue(child, source_child_type, row, doc); if (!child_json) { return tl::make_unexpected(child_json.error()); } @@ -123,8 +141,12 @@ static neug::result valueToJsonValue( case DataTypeId::kArray: { rapidjson::Value arr(rapidjson::kArrayType); const auto& children = ArrayValue::GetChildren(value); + const DataType* source_child_type = nullptr; + if (source_type != nullptr && source_type->id() == DataTypeId::kArray) { + source_child_type = &ArrayType::GetChildType(*source_type); + } for (const auto& child : children) { - auto child_json = valueToJsonValue(child, doc); + auto child_json = valueToJsonValue(child, source_child_type, row, doc); if (!child_json) { return tl::make_unexpected(child_json.error()); } @@ -136,6 +158,10 @@ static neug::result valueToJsonValue( rapidjson::Value obj(rapidjson::kObjectType); const auto& children = StructValue::GetChildren(value); const auto& field_names = StructType::GetFieldNames(value.type()); + const std::vector* source_child_types = nullptr; + if (source_type != nullptr && source_type->id() == DataTypeId::kStruct) { + source_child_types = &StructType::GetChildTypes(*source_type); + } for (size_t i = 0; i < children.size(); ++i) { const auto field_name = i < field_names.size() ? field_names[i] @@ -143,7 +169,12 @@ static neug::result valueToJsonValue( rapidjson::Value key(field_name.c_str(), static_cast(field_name.size()), allocator); - auto child_json = valueToJsonValue(children[i], doc); + const DataType* source_child_type = + source_child_types != nullptr && i < source_child_types->size() + ? &(*source_child_types)[i] + : nullptr; + auto child_json = + valueToJsonValue(children[i], source_child_type, row, doc); if (!child_json) { return tl::make_unexpected(child_json.error()); } @@ -161,36 +192,22 @@ static neug::result valueToJsonValue( } } -static bool isSerializedGraphJson(DataTypeId id) { - return id == DataTypeId::kVertex || id == DataTypeId::kEdge || - id == DataTypeId::kPath; -} - -static const DataType& exportSourceType( +static const DataType* exportSourceType( const std::vector& source_types, size_t col) { - static const DataType kDefaultVarchar(DataTypeId::kVarchar); if (col < source_types.size()) { - return source_types[col]; + return &source_types[col]; } - return kDefaultVarchar; + return nullptr; } -// Convert one cell to a rapidjson value. When source_types[col] is a graph -// type, the VARCHAR payload is pre-serialized JSON and is parsed back into a -// structured value for inline emission. +// Convert one cell to a rapidjson value. Graph leaves in source_type are +// materialized as VARCHAR JSON payloads and parsed back into structured values +// for inline emission, including when nested in a LIST, STRUCT, or ARRAY. static neug::result cellToJsonValue( - const IContextColumn& column, size_t row, const DataType& source_type, + const IContextColumn& column, size_t row, const DataType* source_type, rapidjson::Document& doc) { Value value = column.get_elem(row); - if (isSerializedGraphJson(source_type.id())) { - const auto& str = StringValue::Get(value); - if (str.empty()) { - return rapidjson::Value(rapidjson::kNullType); - } - return parseJsonStringToValue(str, static_cast(row), doc, - source_type.ToString().c_str()); - } - return valueToJsonValue(value, doc); + return valueToJsonValue(value, source_type, static_cast(row), doc); } static Status writeChunkAsJsonArray(const DataChunk& chunk, diff --git a/src/execution/common/operators/retrieve/sink.cc b/src/execution/common/operators/retrieve/sink.cc index b0d0432d7..5cf9ea192 100644 --- a/src/execution/common/operators/retrieve/sink.cc +++ b/src/execution/common/operators/retrieve/sink.cc @@ -615,12 +615,43 @@ std::shared_ptr materialize_graph_column_as_string( return builder->finish(); } +DataType materialized_type_for_export(const DataType& source_type) { + switch (source_type.id()) { + case DataTypeId::kVertex: + case DataTypeId::kEdge: + case DataTypeId::kPath: + return DataType(DataTypeId::kVarchar); + case DataTypeId::kList: + return DataType::List( + materialized_type_for_export(ListType::GetChildType(source_type))); + case DataTypeId::kStruct: { + std::vector materialized_child_types; + const auto& source_child_types = StructType::GetChildTypes(source_type); + materialized_child_types.reserve(source_child_types.size()); + for (const auto& child_type : source_child_types) { + materialized_child_types.push_back( + materialized_type_for_export(child_type)); + } + return DataType::Struct(StructType::GetFieldNames(source_type), + std::move(materialized_child_types)); + } + case DataTypeId::kArray: + return DataType::Array( + materialized_type_for_export(ArrayType::GetChildType(source_type)), + ArrayType::GetNumElements(source_type)); + default: + return source_type; + } +} + Value materialize_value_for_export(const Value& value, + const DataType& source_type, + const DataType& materialized_type, const StorageReadInterface& graph) { if (value.IsNull()) { - return Value(value.type()); + return Value(materialized_type); } - switch (value.type().id()) { + switch (source_type.id()) { case DataTypeId::kVertex: return Value::STRING( convert_vertex_to_json(graph, value.GetValue())); @@ -630,31 +661,43 @@ Value materialize_value_for_export(const Value& value, return Value::STRING(convert_path_to_json(graph, PathValue::Get(value))); case DataTypeId::kList: { const auto& children = ListValue::GetChildren(value); - const auto& child_type = ListType::GetChildType(value.type()); + const auto& source_child_type = ListType::GetChildType(source_type); + const auto& materialized_child_type = + ListType::GetChildType(materialized_type); std::vector materialized; materialized.reserve(children.size()); for (const auto& child : children) { - materialized.push_back(materialize_value_for_export(child, graph)); + materialized.push_back(materialize_value_for_export( + child, source_child_type, materialized_child_type, graph)); } - return Value::LIST(child_type, std::move(materialized)); + return Value::LIST(materialized_child_type, std::move(materialized)); } case DataTypeId::kStruct: { const auto& children = StructValue::GetChildren(value); + const auto& source_child_types = StructType::GetChildTypes(source_type); + const auto& materialized_child_types = + StructType::GetChildTypes(materialized_type); std::vector materialized; materialized.reserve(children.size()); - for (const auto& child : children) { - materialized.push_back(materialize_value_for_export(child, graph)); + for (size_t i = 0; i < children.size(); ++i) { + materialized.push_back( + materialize_value_for_export(children[i], source_child_types[i], + materialized_child_types[i], graph)); } - return Value::STRUCT(value.type(), std::move(materialized)); + return Value::STRUCT(materialized_type, std::move(materialized)); } case DataTypeId::kArray: { const auto& children = ArrayValue::GetChildren(value); + const auto& source_child_type = ArrayType::GetChildType(source_type); + const auto& materialized_child_type = + ArrayType::GetChildType(materialized_type); std::vector materialized; materialized.reserve(children.size()); for (const auto& child : children) { - materialized.push_back(materialize_value_for_export(child, graph)); + materialized.push_back(materialize_value_for_export( + child, source_child_type, materialized_child_type, graph)); } - return Value::ARRAY(value.type(), std::move(materialized)); + return Value::ARRAY(materialized_type, std::move(materialized)); } default: return value; @@ -675,13 +718,16 @@ std::shared_ptr materialize_column_for_export( case DataTypeId::kList: case DataTypeId::kStruct: case DataTypeId::kArray: { - auto builder = ColumnsUtils::create_builder(col->elem_type()); + const auto& source_type = col->elem_type(); + auto materialized_type = materialized_type_for_export(source_type); + auto builder = ColumnsUtils::create_builder(materialized_type); + builder->reserve(col->size()); for (size_t i = 0; i < col->size(); ++i) { if (col->is_optional() && !col->has_value(i)) { builder->push_back_null(); } else { - builder->push_back_elem( - materialize_value_for_export(col->get_elem(i), graph)); + builder->push_back_elem(materialize_value_for_export( + col->get_elem(i), source_type, materialized_type, graph)); } } return builder->finish(); diff --git a/tests/utils/export_result_test.cc b/tests/utils/export_result_test.cc index 9b9430ef8..48ab19b0a 100644 --- a/tests/utils/export_result_test.cc +++ b/tests/utils/export_result_test.cc @@ -34,6 +34,7 @@ #include "neug/storages/graph/graph_view.h" #include "neug/storages/graph/property_graph.h" #include "neug/utils/io/write/writer.h" +#include "unittest/utils.h" namespace neug { namespace test { @@ -174,6 +175,86 @@ TEST_F(ExportResultTest, JsonArrayWriterEmitsNestedValues) { EXPECT_EQ(json, R"([{"payload":{"a":7,"items":[1,2]}}])"); } +TEST_F(ExportResultTest, MaterializerPreservesContainersAroundGraphValues) { + const DataType vertex_type(DataTypeId::kVertex); + const auto vertex_list_type = DataType::List(vertex_type); + const auto vertex_array_type = DataType::Array(vertex_type, 2); + const auto source_type = DataType::Struct( + {"nodes", "primary"}, {vertex_list_type, vertex_array_type}); + + const auto graph_path = + std::string(EXPORT_RESULT_TEST_DIR) + "/nested_graph_data"; + CheckpointManager checkpoint_mgr; + checkpoint_mgr.Open(graph_path); + PropertyGraph graph; + graph.Open(make_checkpoint(checkpoint_mgr), MemoryLevel::kInMemory); + CreateVertexTypeParamBuilder person_builder; + ASSERT_TRUE(graph + .CreateVertexType(person_builder.VertexLabel("person") + .AddProperty("id", Value::INT64(0)) + .AddProperty("name", Value::STRING("")) + .AddPrimaryKeyName("id") + .Build()) + .ok()); + const auto person_label = graph.schema().get_vertex_label_id("person"); + vid_t person_vid; + ASSERT_TRUE(graph + .AddVertex(person_label, Value::INT64(1), + {Value::STRING("Alice")}, person_vid, 0) + .ok()); + const VertexRecord person{person_label, person_vid}; + + std::vector vertices; + vertices.push_back(Value::VERTEX(person)); + vertices.emplace_back(vertex_type); + std::vector primary; + primary.push_back(Value::VERTEX(person)); + primary.emplace_back(vertex_type); + std::vector payload; + payload.push_back(Value::LIST(vertex_type, std::move(vertices))); + payload.push_back(Value::ARRAY(vertex_array_type, std::move(primary))); + + StructColumnBuilder builder(source_type); + builder.push_back_elem(Value::STRUCT(source_type, std::move(payload))); + + execution::Context ctx; + ctx.tag_ids = {0}; + DataChunk input; + input.set(0, builder.finish()); + ctx.append_chunk(std::move(input)); + + GraphView view(graph); + StorageReadInterface reader(view, 1); + auto export_result = materialize_result_for_export(ctx, reader); + + ASSERT_EQ(export_result.source_types.size(), 1); + EXPECT_EQ(export_result.source_types[0], source_type); + const auto materialized_type = + DataType::Struct({"nodes", "primary"}, + {DataType::List(DataType(DataTypeId::kVarchar)), + DataType::Array(DataType(DataTypeId::kVarchar), 2)}); + ASSERT_EQ(export_result.chunk.col_num(), 1); + ASSERT_NE(export_result.chunk.columns[0], nullptr); + EXPECT_EQ(export_result.chunk.columns[0]->elem_type(), materialized_type); + + reader::FileSchema schema; + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/nested_graph.json"}; + schema.format = "json"; + auto entry_schema = std::make_shared(); + entry_schema->columnNames = {"payload"}; + + writer::JsonArrayExportWriter writer(schema, entry_schema); + auto status = writer.write(export_result.chunk, export_result.source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + + std::ifstream file(schema.paths[0]); + std::string json((std::istreambuf_iterator(file)), + std::istreambuf_iterator()); + EXPECT_EQ( + json, + R"([{"payload":{"nodes":[{"_ID":0,"_LABEL":"person","id":1,"name":"Alice"},null],"primary":[{"_ID":0,"_LABEL":"person","id":1,"name":"Alice"},null]}}])"); +} + } // namespace } // namespace test } // namespace neug diff --git a/tools/python_bind/tests/test_export.py b/tools/python_bind/tests/test_export.py index 3874b872d..6be3fb6a8 100644 --- a/tools/python_bind/tests/test_export.py +++ b/tools/python_bind/tests/test_export.py @@ -603,6 +603,20 @@ def test_export_collect_names_jsonl(self): if rows: assert isinstance(rows[0], dict), "Each line should be a JSON object" + def test_export_collect_vertices_jsonl(self): + """Nested graph values remain nested JSON objects after materialization.""" + out_path = self.tmp_path / "collect_vertices.jsonl" + out_path.unlink(missing_ok=True) + self.conn.execute(f"COPY (MATCH (v:person) RETURN collect(v)) TO '{out_path}';") + + rows = _parse_jsonl(out_path) + assert len(rows) == 1 + assert len(rows[0]) == 1 + vertices = next(iter(rows[0].values())) + assert isinstance(vertices, list) + assert vertices + assert all(isinstance(vertex, dict) for vertex in vertices) + @pytest.mark.skipif( not HAS_COMPREHENSIVE_GRAPH, reason="comprehensive_graph data not found" @@ -808,6 +822,22 @@ def test_export_edge_to_parquet(self): # so we only verify the file was created successfully # TODO: Enable LOAD FROM verification when Struct type reading is supported + @extension_test + def test_export_collect_vertices_to_parquet(self): + """Nested graph values are exported as a Parquet list of JSON strings.""" + import pyarrow.parquet as pq + + out_path = self.tmp_path / "collect_vertices.parquet" + self.conn.execute(f"COPY (MATCH (v:person) RETURN collect(v)) TO '{out_path}'") + + table = pq.read_table(out_path) + assert table.num_rows == 1 + assert table.num_columns == 1 + vertices = table.column(0)[0].as_py() + assert isinstance(vertices, list) + assert vertices + assert all(isinstance(json.loads(vertex), dict) for vertex in vertices) + @extension_test def test_export_with_scalar_types(self): """Test Parquet export with various scalar types.""" From 4773a7e1d326d60025101031a366d2e0de7cb653 Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 24 Jul 2026 20:45:05 +0800 Subject: [PATCH 7/8] fix: honor JSON export write options --- src/compiler/function/json_export_function.cc | 83 ++++++++++++++----- tests/utils/export_result_test.cc | 83 +++++++++++++++++++ tools/python_bind/tests/test_export.py | 38 +++++++++ 3 files changed, 185 insertions(+), 19 deletions(-) diff --git a/src/compiler/function/json_export_function.cc b/src/compiler/function/json_export_function.cc index e8770163e..b0c8ec349 100644 --- a/src/compiler/function/json_export_function.cc +++ b/src/compiler/function/json_export_function.cc @@ -73,9 +73,6 @@ static neug::result valueToJsonValue( } if (source_type != nullptr && isSerializedGraphJson(source_type->id())) { const auto& str = StringValue::Get(value); - if (str.empty()) { - return rapidjson::Value(rapidjson::kNullType); - } return parseJsonStringToValue(str, row, doc, source_type->ToString().c_str()); } @@ -213,7 +210,8 @@ static neug::result cellToJsonValue( static Status writeChunkAsJsonArray(const DataChunk& chunk, const reader::FileSchema& schema, const reader::EntrySchema& entry_schema, - const std::vector& source_types) { + const std::vector& source_types, + bool ignore_errors) { if (schema.paths.empty()) { return Status(StatusCode::ERR_INVALID_ARGUMENT, "Schema paths is empty"); } @@ -238,12 +236,23 @@ static Status writeChunkAsJsonArray(const DataChunk& chunk, static_cast(column_name.size()), allocator); if (!chunk.columns[col]->has_value(row)) { + if (!ignore_errors) { + (void) stream->Close(); + return Status(StatusCode::ERR_INVALID_ARGUMENT, + "Value is invalid, rowIdx=" + std::to_string(row) + + ", colIdx=" + std::to_string(col)); + } line.AddMember(key, rapidjson::Value(rapidjson::kNullType), allocator); continue; } auto json_val = cellToJsonValue(*chunk.columns[col], row, exportSourceType(source_types, col), doc); if (!json_val) { + if (ignore_errors) { + line.AddMember(key, rapidjson::Value(rapidjson::kNullType), + allocator); + continue; + } (void) stream->Close(); return json_val.error(); } @@ -265,10 +274,23 @@ static Status writeChunkAsJsonArray(const DataChunk& chunk, return stream->Close(); } +static Status flushJsonLBuffer(io::OutputStream& stream, std::string& buffer) { + if (buffer.empty()) { + return Status::OK(); + } + auto status = stream.Write(reinterpret_cast(buffer.data()), + static_cast(buffer.size())); + if (status.ok()) { + buffer.clear(); + } + return status; +} + static Status writeChunkAsJsonL(const DataChunk& chunk, const reader::FileSchema& schema, const reader::EntrySchema& entry_schema, - const std::vector& source_types) { + const std::vector& source_types, + bool ignore_errors, size_t batch_size) { if (schema.paths.empty()) { return Status(StatusCode::ERR_INVALID_ARGUMENT, "Schema paths is empty"); } @@ -277,6 +299,7 @@ static Status writeChunkAsJsonL(const DataChunk& chunk, return Status(StatusCode::ERR_IO_ERROR, "Failed to open output file"); } + std::string jsonl_buffer; for (size_t row = 0; row < chunk.row_num(); ++row) { rapidjson::Document doc; doc.SetObject(); @@ -292,12 +315,22 @@ static Status writeChunkAsJsonL(const DataChunk& chunk, static_cast(column_name.size()), allocator); if (!chunk.columns[col]->has_value(row)) { + if (!ignore_errors) { + (void) stream->Close(); + return Status(StatusCode::ERR_INVALID_ARGUMENT, + "Value is invalid, rowIdx=" + std::to_string(row) + + ", colIdx=" + std::to_string(col)); + } doc.AddMember(key, rapidjson::Value(rapidjson::kNullType), allocator); continue; } auto json_val = cellToJsonValue(*chunk.columns[col], row, exportSourceType(source_types, col), doc); if (!json_val) { + if (ignore_errors) { + doc.AddMember(key, rapidjson::Value(rapidjson::kNullType), allocator); + continue; + } (void) stream->Close(); return json_val.error(); } @@ -306,20 +339,21 @@ static Status writeChunkAsJsonL(const DataChunk& chunk, rapidjson::StringBuffer buffer; rapidjson::Writer writer(buffer); doc.Accept(writer); - auto status = - stream->Write(reinterpret_cast(buffer.GetString()), - static_cast(buffer.GetSize())); - if (!status.ok()) { - (void) stream->Close(); - return status; - } - status = stream->Write( - reinterpret_cast(DEFAULT_JSON_NEWLINE), sizeof(char)); - if (!status.ok()) { - (void) stream->Close(); - return status; + jsonl_buffer.append(buffer.GetString(), buffer.GetSize()); + jsonl_buffer.append(DEFAULT_JSON_NEWLINE); + if ((row + 1) % batch_size == 0) { + auto status = flushJsonLBuffer(*stream, jsonl_buffer); + if (!status.ok()) { + (void) stream->Close(); + return status; + } } } + auto status = flushJsonLBuffer(*stream, jsonl_buffer); + if (!status.ok()) { + (void) stream->Close(); + return status; + } return stream->Close(); } @@ -334,7 +368,10 @@ Status JsonArrayExportWriter::write(const DataChunk& chunk, std::to_string(chunk.col_num()) + ", got " + std::to_string(source_types.size())); } - return writeChunkAsJsonArray(chunk, schema_, *entry_schema_, source_types); + WriteOptions write_options; + const auto ignore_errors = write_options.ignore_errors.get(schema_.options); + return writeChunkAsJsonArray(chunk, schema_, *entry_schema_, source_types, + ignore_errors); } Status JsonLExportWriter::write(const DataChunk& chunk, @@ -348,7 +385,15 @@ Status JsonLExportWriter::write(const DataChunk& chunk, std::to_string(chunk.col_num()) + ", got " + std::to_string(source_types.size())); } - return writeChunkAsJsonL(chunk, schema_, *entry_schema_, source_types); + WriteOptions write_options; + const auto ignore_errors = write_options.ignore_errors.get(schema_.options); + const auto batch_size = write_options.batch_rows.get(schema_.options); + if (batch_size <= 0) { + return Status(StatusCode::ERR_INVALID_ARGUMENT, + "Batch size should be positive"); + } + return writeChunkAsJsonL(chunk, schema_, *entry_schema_, source_types, + ignore_errors, static_cast(batch_size)); } } // namespace writer diff --git a/tests/utils/export_result_test.cc b/tests/utils/export_result_test.cc index 48ab19b0a..cf7b5e00a 100644 --- a/tests/utils/export_result_test.cc +++ b/tests/utils/export_result_test.cc @@ -70,6 +70,12 @@ std::shared_ptr boolColumn(std::vector values) { return builder.finish(); } +std::shared_ptr nullIntColumn() { + ValueColumnBuilder builder; + builder.push_back_null(); + return builder.finish(); +} + std::vector readLines(const std::string& path) { std::ifstream file(path); std::vector lines; @@ -175,6 +181,83 @@ TEST_F(ExportResultTest, JsonArrayWriterEmitsNestedValues) { EXPECT_EQ(json, R"([{"payload":{"a":7,"items":[1,2]}}])"); } +TEST_F(ExportResultTest, JsonWritersHonorIgnoreErrors) { + reader::FileSchema schema; + schema.format = "json"; + auto entry_schema = std::make_shared(); + entry_schema->columnNames = {"value"}; + + DataChunk null_chunk; + null_chunk.set(0, nullIntColumn()); + const std::vector int_source_types = {DataType(DataTypeId::kInt32)}; + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/strict_null.json"}; + schema.options["IGNORE_ERRORS"] = "false"; + writer::JsonArrayExportWriter strict_null_writer(schema, entry_schema); + auto status = strict_null_writer.write(null_chunk, int_source_types); + EXPECT_FALSE(status.ok()); + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/ignored_null.json"}; + schema.options["IGNORE_ERRORS"] = "true"; + writer::JsonArrayExportWriter ignored_null_writer(schema, entry_schema); + status = ignored_null_writer.write(null_chunk, int_source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + std::ifstream null_file(schema.paths[0]); + std::string null_json((std::istreambuf_iterator(null_file)), + std::istreambuf_iterator()); + EXPECT_EQ(null_json, R"([{"value":null}])"); + + DataChunk invalid_graph_chunk; + invalid_graph_chunk.set(0, stringColumn({"not-json"})); + const std::vector graph_source_types = { + DataType(DataTypeId::kVertex)}; + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + + "/strict_invalid_graph.json"}; + schema.options["IGNORE_ERRORS"] = "false"; + writer::JsonArrayExportWriter strict_graph_writer(schema, entry_schema); + status = strict_graph_writer.write(invalid_graph_chunk, graph_source_types); + EXPECT_FALSE(status.ok()); + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + + "/ignored_invalid_graph.json"}; + schema.options["IGNORE_ERRORS"] = "true"; + writer::JsonArrayExportWriter ignored_graph_writer(schema, entry_schema); + status = ignored_graph_writer.write(invalid_graph_chunk, graph_source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + std::ifstream graph_file(schema.paths[0]); + std::string graph_json((std::istreambuf_iterator(graph_file)), + std::istreambuf_iterator()); + EXPECT_EQ(graph_json, R"([{"value":null}])"); +} + +TEST_F(ExportResultTest, JsonLWriterHonorsBatchSize) { + DataChunk chunk; + chunk.set(0, intColumn({1, 2, 3})); + const std::vector source_types = {DataType(DataTypeId::kInt32)}; + + reader::FileSchema schema; + schema.format = "jsonl"; + auto entry_schema = std::make_shared(); + entry_schema->columnNames = {"id"}; + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/batch_zero.jsonl"}; + schema.options["BATCH_SIZE"] = "0"; + writer::JsonLExportWriter invalid_writer(schema, entry_schema); + auto status = invalid_writer.write(chunk, source_types); + EXPECT_FALSE(status.ok()); + EXPECT_FALSE(std::filesystem::exists(schema.paths[0])); + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/batch_two.jsonl"}; + schema.options["BATCH_SIZE"] = "2"; + writer::JsonLExportWriter writer(schema, entry_schema); + status = writer.write(chunk, source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + EXPECT_EQ( + readLines(schema.paths[0]), + std::vector({R"({"id":1})", R"({"id":2})", R"({"id":3})"})); +} + TEST_F(ExportResultTest, MaterializerPreservesContainersAroundGraphValues) { const DataType vertex_type(DataTypeId::kVertex); const auto vertex_list_type = DataType::List(vertex_type); diff --git a/tools/python_bind/tests/test_export.py b/tools/python_bind/tests/test_export.py index 6be3fb6a8..b61611406 100644 --- a/tools/python_bind/tests/test_export.py +++ b/tools/python_bind/tests/test_export.py @@ -585,6 +585,44 @@ def test_export_person_node_jsonl(self): if rows: assert isinstance(rows[0], dict), "Each line should be a JSON object" + @pytest.mark.parametrize("extension", ["json", "jsonl"]) + def test_export_json_ignore_errors(self, extension): + """JSON writers honor IGNORE_ERRORS for invalid result values.""" + strict_path = self.tmp_path / f"strict_null.{extension}" + with pytest.raises(RuntimeError): + self.conn.execute( + f"COPY (RETURN NULL) TO '{strict_path}' (IGNORE_ERRORS = false);" + ) + + ignored_path = self.tmp_path / f"ignored_null.{extension}" + self.conn.execute( + f"COPY (RETURN NULL) TO '{ignored_path}' (IGNORE_ERRORS = true);" + ) + rows = ( + _parse_json_array(ignored_path) + if extension == "json" + else _parse_jsonl(ignored_path) + ) + assert len(rows) == 1 + assert len(rows[0]) == 1 + assert next(iter(rows[0].values())) is None + + def test_export_jsonl_batch_size(self): + """JSONL validates and uses BATCH_SIZE without changing its rows.""" + invalid_path = self.tmp_path / "batch_zero.jsonl" + with pytest.raises(RuntimeError): + self.conn.execute( + f"COPY (MATCH (v:person) RETURN v.ID) TO '{invalid_path}' " + "(BATCH_SIZE = 0);" + ) + + out_path = self.tmp_path / "batch_two.jsonl" + expected = _count_query(self.conn, "MATCH (v:person) RETURN v.ID") + self.conn.execute( + f"COPY (MATCH (v:person) RETURN v.ID) TO '{out_path}' " "(BATCH_SIZE = 2);" + ) + assert len(_parse_jsonl(out_path)) == expected + def test_export_collect_names_jsonl(self): """Export collect names to JSONL (one JSON object per line); verify row count.""" out_path = self.tmp_path / "collect_names.jsonl" From 1a3a96665e14177ae62afec360af0202c1e0d9ef Mon Sep 17 00:00:00 2001 From: luoxiaojian Date: Fri, 24 Jul 2026 21:24:36 +0800 Subject: [PATCH 8/8] fix: preserve nested values in CSV export --- include/neug/utils/io/write/writer.h | 10 +- src/utils/io/write/writer.cc | 201 +++++++++++++++++++++---- tests/utils/export_result_test.cc | 82 ++++++++++ tools/python_bind/tests/test_export.py | 52 +++++-- 4 files changed, 301 insertions(+), 44 deletions(-) diff --git a/include/neug/utils/io/write/writer.h b/include/neug/utils/io/write/writer.h index cb560e054..29798ee55 100644 --- a/include/neug/utils/io/write/writer.h +++ b/include/neug/utils/io/write/writer.h @@ -81,8 +81,14 @@ class DataChunkCSVStringFormatBuffer { private: neug::Status formatValueToStr(const Value& value, size_t row_idx, size_t col_idx); - bool shouldWriteRawString(size_t col_idx) const; - void writeWithEscapes(char* toEscape, char escape, const std::string& str); + neug::Status appendNestedValue(const Value& value, + const DataType* source_type, + std::string& output) const; + void appendQuotedString(const std::string& value, std::string& output) const; + void writeCsvField(const std::string& value, bool force_quote = false); + const DataType* sourceType(size_t col_idx) const; + void writeWithEscapes(const char* to_escape, char escape, + const std::string& str); void write(const uint8_t* buffer, uint64_t len); const DataChunk& chunk_; diff --git a/src/utils/io/write/writer.cc b/src/utils/io/write/writer.cc index a95700e60..a6589dd7d 100644 --- a/src/utils/io/write/writer.cc +++ b/src/utils/io/write/writer.cc @@ -30,6 +30,18 @@ namespace neug { namespace writer { +namespace { + +bool isSerializedGraphType(const DataType* type) { + if (type == nullptr) { + return false; + } + return type->id() == DataTypeId::kVertex || type->id() == DataTypeId::kEdge || + type->id() == DataTypeId::kPath; +} + +} // namespace + DataChunkCSVStringFormatBuffer::DataChunkCSVStringFormatBuffer( const DataChunk& chunk, const reader::FileSchema& schema, const reader::EntrySchema& entry_schema, @@ -64,14 +76,13 @@ void DataChunkCSVStringFormatBuffer::addHeader() { if (col > 0) { write(reinterpret_cast(&delimiter_), sizeof(char)); } - const auto& name = entry_schema_.columnNames[col]; - write(reinterpret_cast(name.c_str()), name.size()); + writeCsvField(entry_schema_.columnNames[col]); } write(reinterpret_cast(DEFAULT_CSV_NEWLINE), sizeof(char)); } } -void DataChunkCSVStringFormatBuffer::writeWithEscapes(char* to_escape, +void DataChunkCSVStringFormatBuffer::writeWithEscapes(const char* to_escape, char escape, const std::string& val) { uint64_t i = 0; @@ -121,48 +132,166 @@ void DataChunkCSVStringFormatBuffer::write(const uint8_t* buffer, blob_.size += len; } -neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( - const Value& value, size_t row_idx, size_t col_idx) { +void DataChunkCSVStringFormatBuffer::appendQuotedString( + const std::string& value, std::string& output) const { + output.push_back(quote_char_); + for (char c : value) { + if (c == escape_char_ || c == quote_char_) { + output.push_back(escape_char_); + } + output.push_back(c); + } + output.push_back(quote_char_); +} + +const DataType* DataChunkCSVStringFormatBuffer::sourceType( + size_t col_idx) const { + return col_idx < source_types_.size() ? &source_types_[col_idx] : nullptr; +} + +neug::Status DataChunkCSVStringFormatBuffer::appendNestedValue( + const Value& value, const DataType* source_type, + std::string& output) const { if (value.IsNull()) { - return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, - "Value is invalid, rowIdx=" + std::to_string(row_idx)); + output.append("NULL"); + return neug::Status::OK(); + } + if (isSerializedGraphType(source_type)) { + if (value.type().id() != DataTypeId::kVarchar) { + return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, + "Materialized graph value should be VARCHAR, got " + + value.type().ToString()); + } + output.append(StringValue::Get(value)); + return neug::Status::OK(); } + const auto type_id = value.type().id(); - if (type_id == DataTypeId::kVarchar) { - const auto& str = StringValue::Get(value); - if (shouldWriteRawString(col_idx)) { - write(reinterpret_cast(str.c_str()), str.size()); - return neug::Status::OK(); + switch (type_id) { + case DataTypeId::kVarchar: + appendQuotedString(StringValue::Get(value), output); + return neug::Status::OK(); + case DataTypeId::kBoolean: + output.append(value.GetValue() ? "true" : "false"); + return neug::Status::OK(); + case DataTypeId::kList: { + const auto& children = ListValue::GetChildren(value); + const DataType* source_child_type = nullptr; + if (source_type != nullptr && source_type->id() == DataTypeId::kList) { + source_child_type = &ListType::GetChildType(*source_type); + } + output.push_back('['); + for (size_t i = 0; i < children.size(); ++i) { + if (i > 0) { + output.push_back(','); + } + auto status = appendNestedValue(children[i], source_child_type, output); + if (!status.ok()) { + return status; + } + } + output.push_back(']'); + return neug::Status::OK(); + } + case DataTypeId::kArray: { + const auto& children = ArrayValue::GetChildren(value); + const DataType* source_child_type = nullptr; + if (source_type != nullptr && source_type->id() == DataTypeId::kArray) { + source_child_type = &ArrayType::GetChildType(*source_type); + } + output.push_back('['); + for (size_t i = 0; i < children.size(); ++i) { + if (i > 0) { + output.push_back(','); + } + auto status = appendNestedValue(children[i], source_child_type, output); + if (!status.ok()) { + return status; + } } - write(reinterpret_cast("e_char_), sizeof(char)); - char escape_chars[] = {escape_char_, quote_char_}; - writeWithEscapes(escape_chars, escape_char_, str); - write(reinterpret_cast("e_char_), sizeof(char)); + output.push_back(']'); return neug::Status::OK(); } - if (type_id == DataTypeId::kBoolean) { - const auto* str = value.GetValue() ? "true" : "false"; - write(reinterpret_cast(str), strlen(str)); + case DataTypeId::kStruct: { + const auto& children = StructValue::GetChildren(value); + const std::vector* source_child_types = nullptr; + if (source_type != nullptr && source_type->id() == DataTypeId::kStruct) { + source_child_types = &StructType::GetChildTypes(*source_type); + } + output.push_back('['); + for (size_t i = 0; i < children.size(); ++i) { + if (i > 0) { + output.push_back(','); + } + const DataType* source_child_type = + source_child_types != nullptr && i < source_child_types->size() + ? &(*source_child_types)[i] + : nullptr; + auto status = appendNestedValue(children[i], source_child_type, output); + if (!status.ok()) { + return status; + } + } + output.push_back(']'); return neug::Status::OK(); } - const auto& str = value.to_string(); - write(reinterpret_cast(str.c_str()), str.size()); - return neug::Status::OK(); + case DataTypeId::kInt32: + case DataTypeId::kInt64: + case DataTypeId::kUInt32: + case DataTypeId::kUInt64: + case DataTypeId::kFloat: + case DataTypeId::kDouble: + case DataTypeId::kDate: + case DataTypeId::kTimestampMs: + case DataTypeId::kInterval: + output.append(value.to_string()); + return neug::Status::OK(); + default: + return neug::Status( + StatusCode::ERR_NOT_SUPPORTED, + "Unsupported CSV value type: " + value.type().ToString()); + } } -bool DataChunkCSVStringFormatBuffer::shouldWriteRawString( - size_t col_idx) const { - if (col_idx >= source_types_.size()) { - return false; +void DataChunkCSVStringFormatBuffer::writeCsvField(const std::string& value, + bool force_quote) { + const bool needs_quote = force_quote || + value.find(delimiter_) != std::string::npos || + value.find(quote_char_) != std::string::npos || + value.find('\r') != std::string::npos || + value.find('\n') != std::string::npos; + if (!needs_quote) { + write(reinterpret_cast(value.data()), value.size()); + return; } - switch (source_types_[col_idx].id()) { - case DataTypeId::kVertex: - case DataTypeId::kEdge: - case DataTypeId::kPath: - return true; - default: - return false; + + write(reinterpret_cast("e_char_), sizeof(char)); + const char escape_chars[] = {escape_char_, quote_char_, '\0'}; + writeWithEscapes(escape_chars, escape_char_, value); + write(reinterpret_cast("e_char_), sizeof(char)); +} + +neug::Status DataChunkCSVStringFormatBuffer::formatValueToStr( + const Value& value, size_t row_idx, size_t col_idx) { + if (value.IsNull()) { + return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, + "Value is invalid, rowIdx=" + std::to_string(row_idx)); + } + + const auto* source_type = sourceType(col_idx); + if (value.type().id() == DataTypeId::kVarchar && + !isSerializedGraphType(source_type)) { + writeCsvField(StringValue::Get(value), true); + return neug::Status::OK(); + } + + std::string formatted; + auto status = appendNestedValue(value, source_type, formatted); + if (!status.ok()) { + return status; } + writeCsvField(formatted); + return neug::Status::OK(); } void DataChunkCSVStringFormatBuffer::addValue(size_t row_idx, size_t col_idx) { @@ -225,6 +354,12 @@ neug::Status CsvQueryExportWriter::write( return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, "entry_schema is null"); } + if (!source_types.empty() && source_types.size() != chunk.col_num()) { + return neug::Status(StatusCode::ERR_INVALID_ARGUMENT, + "source_types size mismatch: expected " + + std::to_string(chunk.col_num()) + ", got " + + std::to_string(source_types.size())); + } auto stream = io::openLocalOutputStream(schema_.paths[0]); if (!stream) { return neug::Status(StatusCode::ERR_IO_ERROR, "Failed to open output file"); diff --git a/tests/utils/export_result_test.cc b/tests/utils/export_result_test.cc index cf7b5e00a..56593cea1 100644 --- a/tests/utils/export_result_test.cc +++ b/tests/utils/export_result_test.cc @@ -23,9 +23,11 @@ #include #include +#include "neug/common/columns/list_columns.h" #include "neug/common/columns/struct_columns.h" #include "neug/common/columns/value_columns.h" #include "neug/common/export/export_result.h" +#include "neug/common/types/array_columns.h" #include "neug/common/types/data_chunk.h" #include "neug/common/types/value.h" #include "neug/compiler/function/export/json_export_function.h" @@ -144,6 +146,74 @@ TEST_F(ExportResultTest, MaterializerMergesChunksBeforeCsvWrite) { EXPECT_EQ(lines[3], "3|\"carol\"|true"); } +TEST_F(ExportResultTest, CsvWriterFormatsNestedValuesWithoutInformationLoss) { + const DataType string_type(DataTypeId::kVarchar); + ListColumnBuilder string_list_builder(string_type); + std::vector strings; + strings.push_back(Value::STRING("Alice")); + strings.push_back(Value::STRING("a,b")); + strings.push_back(Value::STRING("")); + strings.push_back(Value::STRING("NULL")); + strings.emplace_back(string_type); + string_list_builder.push_back_elem( + Value::LIST(string_type, std::move(strings))); + + const DataType bool_type(DataTypeId::kBoolean); + ListColumnBuilder bool_list_builder(bool_type); + std::vector bools; + bools.push_back(Value::BOOLEAN(true)); + bools.push_back(Value::BOOLEAN(false)); + bool_list_builder.push_back_elem(Value::LIST(bool_type, std::move(bools))); + + const DataType int_type(DataTypeId::kInt32); + const auto int_array_type = DataType::Array(int_type, 2); + ContextArrayColumnBuilder int_array_builder(int_array_type); + std::vector ints; + ints.push_back(Value::INT32(1)); + ints.push_back(Value::INT32(2)); + int_array_builder.push_back_elem( + Value::ARRAY(int_array_type, std::move(ints))); + + const auto struct_type = + DataType::Struct({"id", "name"}, {int_type, string_type}); + StructColumnBuilder struct_builder(struct_type); + std::vector fields; + fields.push_back(Value::INT32(7)); + fields.push_back(Value::STRING("Alice")); + struct_builder.push_back_elem(Value::STRUCT(struct_type, std::move(fields))); + + DataChunk chunk; + chunk.set(0, string_list_builder.finish()); + chunk.set(1, bool_list_builder.finish()); + chunk.set(2, int_array_builder.finish()); + chunk.set(3, struct_builder.finish()); + chunk.set(4, stringColumn({R"({"id":1,"name":"A,B"})"})); + + reader::FileSchema schema; + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/nested.csv"}; + schema.format = "csv"; + schema.options["HEADER"] = "false"; + auto entry_schema = std::make_shared(); + entry_schema->columnNames = {"strings", "bools", "ints", "record", "graph"}; + + writer::CsvQueryExportWriter writer(schema, entry_schema); + const std::vector source_types = { + DataType::List(string_type), DataType::List(bool_type), int_array_type, + struct_type, DataType(DataTypeId::kVertex)}; + auto status = writer.write(chunk, {DataType::List(string_type)}); + EXPECT_FALSE(status.ok()); + EXPECT_FALSE(std::filesystem::exists(schema.paths[0])); + + status = writer.write(chunk, source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + + auto lines = readLines(schema.paths[0]); + ASSERT_EQ(lines.size(), 1); + EXPECT_EQ( + lines[0], + R"("[\"Alice\",\"a,b\",\"\",\"NULL\",NULL]"|[true,false]|[1,2]|"[7,\"Alice\"]"|"{\"id\":1,\"name\":\"A,B\"}")"); +} + TEST_F(ExportResultTest, JsonArrayWriterEmitsNestedValues) { auto list_type = DataType::List(DataType(DataTypeId::kInt32)); std::vector field_names = {"a", "items"}; @@ -336,6 +406,18 @@ TEST_F(ExportResultTest, MaterializerPreservesContainersAroundGraphValues) { EXPECT_EQ( json, R"([{"payload":{"nodes":[{"_ID":0,"_LABEL":"person","id":1,"name":"Alice"},null],"primary":[{"_ID":0,"_LABEL":"person","id":1,"name":"Alice"},null]}}])"); + + schema.paths = {std::string(EXPORT_RESULT_TEST_DIR) + "/nested_graph.csv"}; + schema.format = "csv"; + schema.options["HEADER"] = "false"; + writer::CsvQueryExportWriter csv_writer(schema, entry_schema); + status = csv_writer.write(export_result.chunk, export_result.source_types); + ASSERT_TRUE(status.ok()) << status.ToString(); + + EXPECT_EQ( + readLines(schema.paths[0]), + std::vector( + {R"("[[{\"_ID\":0,\"_LABEL\":\"person\",\"id\":1,\"name\":\"Alice\"},NULL],[{\"_ID\":0,\"_LABEL\":\"person\",\"id\":1,\"name\":\"Alice\"},NULL]]")"})); } } // namespace diff --git a/tools/python_bind/tests/test_export.py b/tools/python_bind/tests/test_export.py index b61611406..2169d1e68 100644 --- a/tools/python_bind/tests/test_export.py +++ b/tools/python_bind/tests/test_export.py @@ -16,6 +16,7 @@ # limitations under the License. # +import ast import csv import json import os @@ -44,10 +45,21 @@ def _count_query(conn, cypher): return len(list(conn.execute(cypher))) -def _parse_csv(path, delimiter="|", has_header=True): +def _parse_csv( + path, + delimiter="|", + has_header=True, + quotechar='"', + escapechar="\\", +): """Parse CSV; returns (header or None, list of data rows).""" with open(path, newline="", encoding="utf-8") as f: - reader = csv.reader(f, delimiter=delimiter) + reader = csv.reader( + f, + delimiter=delimiter, + quotechar=quotechar, + escapechar=escapechar, + ) rows = list(reader) if not rows: return (None, []) @@ -464,17 +476,39 @@ def test_export_combined_options(self): def test_export_collect_names(self): out_path = self.tmp_path / "collect_names.csv" out_path.unlink(missing_ok=True) - expected = _count_query( - self.conn, "MATCH (v:person) RETURN v.ID, collect(v.fName)" - ) + query = "MATCH (v:person) RETURN v.ID, collect(v.fName)" + expected = {str(row[0]): sorted(row[1]) for row in self.conn.execute(query)} self.conn.execute( - f"COPY (MATCH (v:person) RETURN v.ID, collect(v.fName)) TO " - f"'{out_path}' (HEADER = true, QUOTE = '\\'');" + f"COPY ({query}) TO " f"'{out_path}' (HEADER = true, QUOTE = '\\'');" ) assert out_path.exists() - header, rows = _parse_csv(out_path, "|", has_header=True) + header, rows = _parse_csv( + out_path, + "|", + has_header=True, + quotechar="'", + ) assert len(header) == 2 - assert len(rows) == expected + assert len(rows) == len(expected) + for row in rows: + assert len(row) == 2 + assert sorted(ast.literal_eval(row[1])) == expected[row[0]] + + def test_export_nested_strings_with_comma_delimiter(self): + out_path = self.tmp_path / "nested_comma.csv" + query = "MATCH (v:person) RETURN v.gender, collect(v.fName)" + expected = {str(row[0]): sorted(row[1]) for row in self.conn.execute(query)} + + self.conn.execute( + f"COPY ({query}) TO '{out_path}' " "(HEADER = true, DELIMITER = ',');" + ) + + header, rows = _parse_csv(out_path, ",", has_header=True) + assert len(header) == 2 + assert len(rows) == len(expected) + for row in rows: + assert len(row) == 2 + assert sorted(json.loads(row[1])) == expected[row[0]] # Verify that the 'QUOTE' option correctly changes the wrapping character for string values. # Here, we explicitly set QUOTE = "'" (single quote).