test

Author: sollhui

be/src/common/config.cpp

@@ -1680,6 +1680,10 @@ DEFINE_mBool(enable_auto_clone_on_mow_publish_missing_version, "false");
 // The maximum csv line reader output buffer size
 DEFINE_mInt64(max_csv_line_reader_output_buffer_size, "4294967296");
 
+// The maximum bytes of a single block returned by CsvReader::get_next_block.
+// Default is 200MB. Set to 0 to disable the limit.
+DEFINE_mInt64(csv_reader_max_block_bytes, "209715200");
+
 // Maximum number of OpenMP threads allowed for concurrent vector index builds.
 // -1 means auto: use 80% of the available CPU cores.
 DEFINE_Int32(omp_threads_limit, "-1");

be/src/common/config.h

@@ -1770,6 +1770,10 @@ DECLARE_String(fuzzy_test_type);
 // The maximum csv line reader output buffer size
 DECLARE_mInt64(max_csv_line_reader_output_buffer_size);
 
+// The maximum bytes of a single block returned by CsvReader::get_next_block.
+// Default is 200MB. Set to 0 to disable the limit.
+DECLARE_mInt64(csv_reader_max_block_bytes);
+
 // Maximum number of OpenMP threads available for concurrent index builds.
 // -1 means auto: use 80% of detected CPU cores.
 DECLARE_Int32(omp_threads_limit);

be/src/format/csv/csv_reader.cpp

@@ -31,6 +31,7 @@
 #include <utility>
 
 #include "common/compiler_util.h" // IWYU pragma: keep
+#include "common/config.h"
 #include "common/consts.h"
 #include "common/status.h"
 #include "core/block/block.h"
@@ -312,12 +313,15 @@ Status CsvReader::get_next_block(Block* block, size_t* read_rows, bool* eof) {
     }
 
     const int batch_size = std::max(_state->batch_size(), (int)_MIN_BATCH_SIZE);
+    const int64_t max_block_bytes = config::csv_reader_max_block_bytes;
     size_t rows = 0;
+    size_t block_bytes = 0;
 
     bool success = false;
     bool is_remove_bom = false;
     if (_push_down_agg_type == TPushAggOp::type::COUNT) {
-        while (rows < batch_size && !_line_reader_eof) {
+        while (rows < batch_size && !_line_reader_eof &&
+               (max_block_bytes <= 0 || (int64_t)block_bytes < max_block_bytes)) {
             const uint8_t* ptr = nullptr;
             size_t size = 0;
             RETURN_IF_ERROR(_line_reader->read_line(&ptr, &size, &_line_reader_eof, _io_ctx));
@@ -345,6 +349,7 @@ Status CsvReader::get_next_block(Block* block, size_t* read_rows, bool* eof) {
 
             RETURN_IF_ERROR(_validate_line(Slice(ptr, size), &success));
             ++rows;
+            block_bytes += size;
         }
         auto mutate_columns = block->mutate_columns();
         for (auto& col : mutate_columns) {
@@ -353,7 +358,8 @@ Status CsvReader::get_next_block(Block* block, size_t* read_rows, bool* eof) {
         block->set_columns(std::move(mutate_columns));
     } else {
         auto columns = block->mutate_columns();
-        while (rows < batch_size && !_line_reader_eof) {
+        while (rows < batch_size && !_line_reader_eof &&
+               (max_block_bytes <= 0 || (int64_t)block_bytes < max_block_bytes)) {
             const uint8_t* ptr = nullptr;
             size_t size = 0;
             RETURN_IF_ERROR(_line_reader->read_line(&ptr, &size, &_line_reader_eof, _io_ctx));
@@ -384,6 +390,7 @@ Status CsvReader::get_next_block(Block* block, size_t* read_rows, bool* eof) {
                 continue;
             }
             RETURN_IF_ERROR(_fill_dest_columns(Slice(ptr, size), block, columns, &rows));
+            block_bytes += size;
         }
         block->set_columns(std::move(columns));
     }