Enable Partition Discovery for Broker Load

be/src/env/env_posix.cpp

@@ -15,6 +15,7 @@
 #include <sys/stat.h>
 #include <sys/types.h>
 #include <unistd.h>
+#include <sys/uio.h>
 
 #include "common/logging.h"
 #include "gutil/macros.h"

be/src/exec/broker_scanner.cpp

@@ -54,6 +54,7 @@ BrokerScanner::BrokerScanner(RuntimeState* state,
         _cur_decompressor(nullptr),
         _next_range(0),
         _cur_line_reader_eof(false),
+        _columns_from_path(),
         _scanner_eof(false),
         _skip_next_line(false) {
 }
@@ -73,7 +74,7 @@ Status BrokerScanner::open() {
 
 Status BrokerScanner::get_next(Tuple* tuple, MemPool* tuple_pool, bool* eof) {
     SCOPED_TIMER(_read_timer);
-    // Get one line 
+    // Get one line
     while (!_scanner_eof) {
         if (_cur_line_reader == nullptr || _cur_line_reader_eof) {
             RETURN_IF_ERROR(open_next_reader());
@@ -120,7 +121,7 @@ Status BrokerScanner::open_next_reader() {
     RETURN_IF_ERROR(open_file_reader());
     RETURN_IF_ERROR(open_line_reader());
     _next_range++;
-    
+
     return Status::OK();
 }
 
@@ -236,6 +237,9 @@ Status BrokerScanner::open_line_reader() {
     // _decompressor may be NULL if this is not a compressed file
     RETURN_IF_ERROR(create_decompressor(range.format_type));
 
+    // set partitioned columns
+    _columns_from_path = range.columns_from_path;
+
     // open line reader
     switch (range.format_type) {
     case TFileFormatType::FORMAT_CSV_PLAIN:
@@ -379,8 +383,8 @@ bool BrokerScanner::check_decimal_input(
 }
 
 bool is_null(const Slice& slice) {
-    return slice.size == 2 && 
-        slice.data[0] == '\\' && 
+    return slice.size == 2 &&
+        slice.data[0] == '\\' &&
         slice.data[1] == 'N';
 }
 
@@ -452,6 +456,17 @@ bool BrokerScanner::convert_one_row(
     return fill_dest_tuple(line, tuple, tuple_pool);
 }
 
+inline void BrokerScanner::fill_slot(SlotDescriptor* slot_desc, const Slice& value) {
+    if (slot_desc->is_nullable() && is_null(value)) {
+        _src_tuple->set_null(slot_desc->null_indicator_offset());
+    }
+    _src_tuple->set_not_null(slot_desc->null_indicator_offset());
+    void* slot = _src_tuple->get_slot(slot_desc->tuple_offset());
+    StringValue* str_slot = reinterpret_cast<StringValue*>(slot);
+    str_slot->ptr = value.data;
+    str_slot->len = value.size;
+}
+
 // Convert one row to this tuple
 bool BrokerScanner::line_to_src_tuple(const Slice& line) {
 
@@ -469,7 +484,7 @@ bool BrokerScanner::line_to_src_tuple(const Slice& line) {
         split_line(line, &values);
     }
 
-    if (values.size() < _src_slot_descs.size()) {
+    if (values.size() + _columns_from_path.size() < _src_slot_descs.size()) {
         std::stringstream error_msg;
         error_msg << "actual column number is less than schema column number. "
             << "actual number: " << values.size() << " sep: " << _value_separator << ", "
@@ -478,7 +493,7 @@ bool BrokerScanner::line_to_src_tuple(const Slice& line) {
                                          error_msg.str());
         _counter->num_rows_filtered++;
         return false;
-    } else if (values.size() > _src_slot_descs.size()) {
+    } else if (values.size() + _columns_from_path.size() > _src_slot_descs.size()) {
         std::stringstream error_msg;
         error_msg << "actual column number is more than schema column number. "
             << "actual number: " << values.size() << " sep: " << _value_separator << ", "
@@ -489,18 +504,18 @@ bool BrokerScanner::line_to_src_tuple(const Slice& line) {
         return false;
     }
 
-    for (int i = 0; i < values.size(); ++i) {
+    int file_column_index = 0;
+    for (int i = 0; i < _src_slot_descs.size(); ++i) {
         auto slot_desc = _src_slot_descs[i];
-        const Slice& value = values[i];
-        if (slot_desc->is_nullable() && is_null(value)) {
-            _src_tuple->set_null(slot_desc->null_indicator_offset());
-            continue;
+        auto iter = _columns_from_path.find(slot_desc->col_name());
+        if (iter != _columns_from_path.end()) {
+            std::string partitioned_field = iter->second;
+            const Slice value = Slice(partitioned_field.c_str(), partitioned_field.size());
+            fill_slot(slot_desc, value);
+        } else {
+            const Slice& value = values[file_column_index++];
+            fill_slot(slot_desc, value);
         }
-        _src_tuple->set_not_null(slot_desc->null_indicator_offset());
-        void* slot = _src_tuple->get_slot(slot_desc->tuple_offset());
-        StringValue* str_slot = reinterpret_cast<StringValue*>(slot);
-        str_slot->ptr = value.data;
-        str_slot->len = value.size;
     }
 
     return true;

be/src/exec/broker_scanner.h

@@ -104,6 +104,7 @@ class BrokerScanner : public BaseScanner {
     //Status init_expr_ctxes();
 
     Status line_to_src_tuple();
+    void fill_slot(SlotDescriptor* slot_desc, const Slice& value);
     bool line_to_src_tuple(const Slice& line);
 private:;
     const std::vector<TBrokerRangeDesc>& _ranges;
@@ -120,6 +121,7 @@ private:;
     Decompressor* _cur_decompressor;
     int _next_range;
     bool _cur_line_reader_eof;
+    std::map<std::string, std::string> _columns_from_path;
 
     bool _scanner_eof;