diff --git a/src/silo/query_engine/exec_node/ndjson_sink.cpp b/src/silo/query_engine/exec_node/ndjson_sink.cpp index dd3740c5b..eb7f0153f 100644 --- a/src/silo/query_engine/exec_node/ndjson_sink.cpp +++ b/src/silo/query_engine/exec_node/ndjson_sink.cpp @@ -1,5 +1,8 @@ #include "silo/query_engine/exec_node/ndjson_sink.h" +#include +#include + #include #include #include @@ -8,41 +11,53 @@ #include #include "evobench/evobench.hpp" +#include "silo/common/panic.h" #include "silo/common/size_constants.h" namespace silo::query_engine::exec_node { namespace { -template +void writeChunked(std::ostream& output, std::string_view content) { + const size_t chunk_size = 8192; + for (size_t pos = 0; pos < content.size(); pos += chunk_size) { + size_t remaining_size = content.size() - pos; + size_t write_size = std::min(chunk_size, remaining_size); + output.write(content.data() + pos, static_cast(write_size)); + output.flush(); // Flush after each small chunk + } +} + +template struct BatchedStringStream { - std::array streams; + std::array streams; void operator<<(std::string_view bytes) { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { streams[i] << bytes; } } void operator>>(std::ostream& output) { - for (size_t i = 0; i < BATCH_SIZE; ++i) { - output << streams[i].rdbuf(); + for (size_t i = 0; i < BatchSize; ++i) { + std::string content = std::move(streams[i]).str(); + writeChunked(output, content); } } }; -template +template class ArrayToJsonTypeVisitor : public arrow::ArrayVisitor { - BatchedStringStream& output_stream; + BatchedStringStream& output_stream; size_t& row_base; public: - ArrayToJsonTypeVisitor(BatchedStringStream& output_stream, size_t& row_base) + ArrayToJsonTypeVisitor(BatchedStringStream& output_stream, size_t& row_base) : output_stream(output_stream), row_base(row_base) {} arrow::Status Visit(const arrow::Int32Array& array) override { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { if (array.IsNull(row_base + i)) { output_stream.streams.at(i) << "null"; } else { @@ -53,7 +68,7 @@ class ArrayToJsonTypeVisitor : public arrow::ArrayVisitor { } arrow::Status Visit(const arrow::Int64Array& array) override { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { if (array.IsNull(row_base + i)) { output_stream.streams.at(i) << "null"; } else { @@ -64,31 +79,31 @@ class ArrayToJsonTypeVisitor : public arrow::ArrayVisitor { } arrow::Status Visit(const arrow::DoubleArray& array) override { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { if (array.IsNull(row_base + i)) { output_stream.streams.at(i) << "null"; } else { - nlohmann::json j = array.GetView(row_base + i); - output_stream.streams.at(i) << j; + nlohmann::json json = array.GetView(row_base + i); + output_stream.streams.at(i) << json; } } return arrow::Status::OK(); } arrow::Status Visit(const arrow::FloatArray& array) override { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { if (array.IsNull(row_base + i)) { output_stream.streams.at(i) << "null"; } else { - nlohmann::json j = array.GetView(row_base + i); - output_stream.streams.at(i) << j; + nlohmann::json json = array.GetView(row_base + i); + output_stream.streams.at(i) << json; } } return arrow::Status::OK(); } arrow::Status Visit(const arrow::StringArray& array) override { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { if (array.IsNull(row_base + i)) { output_stream.streams.at(i) << "null"; } else { @@ -101,7 +116,7 @@ class ArrayToJsonTypeVisitor : public arrow::ArrayVisitor { } arrow::Status Visit(const arrow::BooleanArray& array) override { - for (size_t i = 0; i < BATCH_SIZE; ++i) { + for (size_t i = 0; i < BatchSize; ++i) { if (array.IsNull(row_base + i)) { output_stream.streams.at(i) << "null"; } else { @@ -112,7 +127,7 @@ class ArrayToJsonTypeVisitor : public arrow::ArrayVisitor { } }; -template +template void sendJsonLinesInBatches( size_t& row_idx_base, size_t row_count, @@ -120,10 +135,10 @@ void sendJsonLinesInBatches( const std::vector>& column_arrays, std::ostream& output_stream ) { - BatchedStringStream ndjson_line_streams; - ArrayToJsonTypeVisitor my_visitor(ndjson_line_streams, row_idx_base); + BatchedStringStream ndjson_line_streams; + ArrayToJsonTypeVisitor my_visitor(ndjson_line_streams, row_idx_base); size_t column_count = column_arrays.size(); - for (; row_idx_base + BATCH_SIZE <= row_count; row_idx_base += BATCH_SIZE) { + for (; row_idx_base + BatchSize <= row_count; row_idx_base += BatchSize) { ndjson_line_streams << "{"; for (size_t column_idx = 0; column_idx < column_count; column_idx++) { const auto& column_array = column_arrays.at(column_idx); @@ -137,7 +152,7 @@ void sendJsonLinesInBatches( ndjson_line_streams >> output_stream; } } - if constexpr (BATCH_SIZE > 1) { + if constexpr (BatchSize > 1) { // Send remaining lines sendJsonLinesInBatches<1>( row_idx_base, @@ -152,7 +167,7 @@ void sendJsonLinesInBatches( } // namespace arrow::Status writeBatchAsNdjson( - arrow::compute::ExecBatch batch, + const arrow::compute::ExecBatch& batch, const std::shared_ptr& schema, std::ostream* output_stream ) { @@ -168,8 +183,9 @@ arrow::Status writeBatchAsNdjson( for (const auto& column_name : schema->fields()) { nlohmann::json column_name_json = column_name->name(); std::string json_formatted_column_name; - if (!first_column) + if (!first_column) { json_formatted_column_name += ","; + } first_column = false; json_formatted_column_name += column_name_json.dump(); json_formatted_column_name += ":"; diff --git a/src/silo/query_engine/exec_node/ndjson_sink.h b/src/silo/query_engine/exec_node/ndjson_sink.h index 3ffd650ff..85f14ef16 100644 --- a/src/silo/query_engine/exec_node/ndjson_sink.h +++ b/src/silo/query_engine/exec_node/ndjson_sink.h @@ -6,12 +6,10 @@ #include #include -#include "silo/common/panic.h" - namespace silo::query_engine::exec_node { arrow::Status writeBatchAsNdjson( - arrow::compute::ExecBatch batch, + const arrow::compute::ExecBatch& batch, const std::shared_ptr& schema, std::ostream* output_stream );