diff --git a/native-sql-engine/cpp/src/benchmarks/shuffle_split_benchmark.cc b/native-sql-engine/cpp/src/benchmarks/shuffle_split_benchmark.cc index 630ad8905..c1dc21aeb 100644 --- a/native-sql-engine/cpp/src/benchmarks/shuffle_split_benchmark.cc +++ b/native-sql-engine/cpp/src/benchmarks/shuffle_split_benchmark.cc @@ -29,6 +29,7 @@ #include #include #include +#include #include #include @@ -159,7 +160,13 @@ class LargePageMemoryPool : public MemoryPool { class BenchmarkShuffleSplit { public: - BenchmarkShuffleSplit(std::string file_name) { GetRecordBatchReader(file_name); } + BenchmarkShuffleSplit(const std::string& file_name, + const std::string& spark_local_dirs = "") { + GetRecordBatchReader(file_name); + setenv("NATIVESQL_SPARK_LOCAL_DIRS", spark_local_dirs.c_str(), 1); + + ARROW_ASSIGN_OR_THROW(spark_local_dirs_, GetConfiguredLocalDirs()); + } void GetRecordBatchReader(const std::string& input_file) { std::unique_ptr<::parquet::arrow::FileReader> parquet_reader; @@ -213,19 +220,24 @@ class BenchmarkShuffleSplit { options.compression_type = compression_type; options.buffer_size = split_buffer_size; options.buffered_write = true; - options.offheap_per_task = 128 * 1024 * 1024 * 1024L; + options.offheap_per_task = 7 * 1024 * 1024 * 1024L; options.prefer_spill = true; options.write_schema = false; options.memory_pool = pool.get(); + options.data_file = arrow::fs::internal::ConcatAbstractPath( + spark_local_dirs_[state.thread_index() % spark_local_dirs_.size()], + "temp_shuffle_" + GenerateUUID()); + std::shared_ptr splitter; int64_t elapse_read = 0; int64_t num_batches = 0; int64_t num_rows = 0; int64_t split_time = 0; + int64_t dur_time = 0; - Do_Split(splitter, elapse_read, num_batches, num_rows, split_time, num_partitions, - options, state); + Do_Split(splitter, elapse_read, num_batches, num_rows, split_time, dur_time, + num_partitions, options, state); auto fs = std::make_shared(); fs->DeleteFile(splitter->DataFile()); @@ -285,6 +297,8 @@ class BenchmarkShuffleSplit { splitter->TotalWriteTime(); state.counters["split_time"] = benchmark::Counter( split_time, benchmark::Counter::kAvgThreads, benchmark::Counter::OneK::kIs1000); + state.counters["elapsed_time"] = benchmark::Counter( + dur_time, benchmark::Counter::kAvgThreads, benchmark::Counter::OneK::kIs1000); splitter.reset(); } @@ -297,7 +311,7 @@ class BenchmarkShuffleSplit { } virtual void Do_Split(std::shared_ptr& splitter, int64_t& elapse_read, int64_t& num_batches, int64_t& num_rows, int64_t& split_time, - const int num_partitions, SplitOptions options, + int64_t& dur_time, const int num_partitions, SplitOptions options, benchmark::State& state) {} protected: @@ -308,39 +322,42 @@ class BenchmarkShuffleSplit { std::shared_ptr schema; std::vector> expr_vector; parquet::ArrowReaderProperties properties; + std::vector spark_local_dirs_; }; class BenchmarkShuffleSplit_CacheScan_Benchmark : public BenchmarkShuffleSplit { public: - BenchmarkShuffleSplit_CacheScan_Benchmark(std::string filename) - : BenchmarkShuffleSplit(filename) {} + BenchmarkShuffleSplit_CacheScan_Benchmark(std::string filename, + const std::string& spark_local_dirs = "") + : BenchmarkShuffleSplit(filename, spark_local_dirs) {} protected: void Do_Split(std::shared_ptr& splitter, int64_t& elapse_read, int64_t& num_batches, int64_t& num_rows, int64_t& split_time, - const int num_partitions, SplitOptions options, benchmark::State& state) { + int64_t& dur_time, const int num_partitions, SplitOptions options, + benchmark::State& state) { std::vector local_column_indices; - // local_column_indices.push_back(0); - /* local_column_indices.push_back(0); - local_column_indices.push_back(1); - local_column_indices.push_back(2); - local_column_indices.push_back(4); - local_column_indices.push_back(5); - local_column_indices.push_back(6); - local_column_indices.push_back(7); -*/ + + local_column_indices.push_back(0); + local_column_indices.push_back(1); + local_column_indices.push_back(2); + local_column_indices.push_back(4); + local_column_indices.push_back(5); + local_column_indices.push_back(6); local_column_indices.push_back(8); local_column_indices.push_back(9); - local_column_indices.push_back(13); - local_column_indices.push_back(14); local_column_indices.push_back(15); std::shared_ptr local_schema; arrow::FieldVector fields; + fields.push_back(schema->field(0)); + fields.push_back(schema->field(1)); + fields.push_back(schema->field(2)); + fields.push_back(schema->field(4)); + fields.push_back(schema->field(5)); + fields.push_back(schema->field(6)); fields.push_back(schema->field(8)); fields.push_back(schema->field(9)); - fields.push_back(schema->field(13)); - fields.push_back(schema->field(14)); fields.push_back(schema->field(15)); local_schema = std::make_shared(fields); @@ -360,6 +377,8 @@ class BenchmarkShuffleSplit_CacheScan_Benchmark : public BenchmarkShuffleSplit { std::vector> batches; ASSERT_NOT_OK(parquet_reader->GetRecordBatchReader( row_group_indices, local_column_indices, &record_batch_reader)); + + auto start_time = std::chrono::steady_clock::now(); do { TIME_NANO_OR_THROW(elapse_read, record_batch_reader->ReadNext(&record_batch)); @@ -384,26 +403,52 @@ class BenchmarkShuffleSplit_CacheScan_Benchmark : public BenchmarkShuffleSplit { } TIME_NANO_OR_THROW(split_time, splitter->Stop()); + auto end_time = std::chrono::steady_clock::now(); + dur_time = std::chrono::duration_cast(end_time - start_time) + .count(); } }; class BenchmarkShuffleSplit_IterateScan_Benchmark : public BenchmarkShuffleSplit { public: - BenchmarkShuffleSplit_IterateScan_Benchmark(std::string filename) - : BenchmarkShuffleSplit(filename) {} + BenchmarkShuffleSplit_IterateScan_Benchmark(std::string filename, + const std::string& spark_local_dirs = "") + : BenchmarkShuffleSplit(filename, spark_local_dirs) {} protected: void Do_Split(std::shared_ptr& splitter, int64_t& elapse_read, int64_t& num_batches, int64_t& num_rows, int64_t& split_time, - const int num_partitions, SplitOptions options, benchmark::State& state) { - if (state.thread_index() == 0) std::cout << schema->ToString() << std::endl; + int64_t& dur_time, const int num_partitions, SplitOptions options, + benchmark::State& state) { + // if (state.thread_index() == 0) std::cout << schema->ToString() << std::endl; + + std::vector local_column_indices; + + local_column_indices.push_back(0); + local_column_indices.push_back(1); + local_column_indices.push_back(2); + local_column_indices.push_back(4); + local_column_indices.push_back(5); + local_column_indices.push_back(6); + + std::shared_ptr local_schema; + arrow::FieldVector fields; + fields.push_back(schema->field(0)); + fields.push_back(schema->field(1)); + fields.push_back(schema->field(2)); + fields.push_back(schema->field(4)); + fields.push_back(schema->field(5)); + fields.push_back(schema->field(6)); + local_schema = std::make_shared(fields); + + if (state.thread_index() == 0) std::cout << local_schema->ToString() << std::endl; if (!expr_vector.empty()) { - ARROW_ASSIGN_OR_THROW(splitter, Splitter::Make("hash", schema, num_partitions, + ARROW_ASSIGN_OR_THROW(splitter, Splitter::Make("hash", local_schema, num_partitions, expr_vector, std::move(options))); } else { - ARROW_ASSIGN_OR_THROW( - splitter, Splitter::Make("rr", schema, num_partitions, std::move(options))); + ARROW_ASSIGN_OR_THROW(splitter, Splitter::Make("rr", local_schema, num_partitions, + std::move(options))); } std::shared_ptr record_batch; @@ -414,10 +459,11 @@ class BenchmarkShuffleSplit_IterateScan_Benchmark : public BenchmarkShuffleSplit arrow::default_memory_pool(), ::parquet::ParquetFileReader::Open(file), properties, &parquet_reader)); + auto start_time = std::chrono::steady_clock::now(); for (auto _ : state) { std::vector> batches; ASSERT_NOT_OK(parquet_reader->GetRecordBatchReader( - row_group_indices, column_indices, &record_batch_reader)); + row_group_indices, local_column_indices, &record_batch_reader)); TIME_NANO_OR_THROW(elapse_read, record_batch_reader->ReadNext(&record_batch)); while (record_batch) { num_batches += 1; @@ -427,6 +473,9 @@ class BenchmarkShuffleSplit_IterateScan_Benchmark : public BenchmarkShuffleSplit } } TIME_NANO_OR_THROW(split_time, splitter->Stop()); + auto end_time = std::chrono::steady_clock::now(); + dur_time = std::chrono::duration_cast(end_time - start_time) + .count(); } }; @@ -472,6 +521,7 @@ int main(int argc, char** argv) { uint32_t partitions = 512; uint32_t threads = 1; std::string datafile; + std::string spark_local_dirs = ""; for (int i = 0; i < argc; i++) { if (strcmp(argv[i], "--iterations") == 0) { @@ -482,6 +532,8 @@ int main(int argc, char** argv) { threads = atol(argv[i + 1]); } else if (strcmp(argv[i], "--file") == 0) { datafile = argv[i + 1]; + } else if (strcmp(argv[i], "--local_dirs") == 0) { + spark_local_dirs = argv[i + 1]; } } std::cout << "iterations = " << iterations << std::endl; @@ -489,41 +541,43 @@ int main(int argc, char** argv) { std::cout << "threads = " << threads << std::endl; std::cout << "datafile = " << datafile << std::endl; - sparkcolumnarplugin::shuffle::BenchmarkShuffleSplit_CacheScan_Benchmark bck(datafile); + /* sparkcolumnarplugin::shuffle::BenchmarkShuffleSplit_CacheScan_Benchmark + bck(datafile, spark_local_dirs); + benchmark::RegisterBenchmark("BenchmarkShuffleSplit::CacheScan", bck) + ->Iterations(iterations) + ->Args({partitions, arrow::Compression::FASTPFOR}) + ->Threads(threads) + ->ReportAggregatesOnly(false) + ->MeasureProcessCPUTime() + ->Unit(benchmark::kSecond); + */ + sparkcolumnarplugin::shuffle::BenchmarkShuffleSplit_IterateScan_Benchmark bck( + datafile, spark_local_dirs); - benchmark::RegisterBenchmark("BenchmarkShuffleSplit::CacheScan", bck) + benchmark::RegisterBenchmark("BenchmarkShuffleSplit::IterateScan", bck) ->Iterations(iterations) ->Args({partitions, arrow::Compression::FASTPFOR}) ->Threads(threads) - ->ReportAggregatesOnly(false) + ->Unit(benchmark::kSecond) ->MeasureProcessCPUTime() - ->Unit(benchmark::kSecond); - - /* sparkcolumnarplugin::shuffle::BenchmarkShuffleSplit_IterateScan_Benchmark - bck(datafile); - - benchmark::RegisterBenchmark("BenchmarkShuffleSplit::IterateScan", bck) - ->Iterations(1) - ->Args({96*2, arrow::Compression::FASTPFOR}) - ->Args({96*4, arrow::Compression::FASTPFOR}) - ->Args({96*8, arrow::Compression::FASTPFOR}) - ->Args({96*16, arrow::Compression::FASTPFOR}) - ->Args({96*32, arrow::Compression::FASTPFOR}) - ->Threads(24) - ->Unit(benchmark::kSecond); - - benchmark::RegisterBenchmark("BenchmarkShuffleSplit::IterateScan", bck) - ->Iterations(1) - ->Args({4096, arrow::Compression::FASTPFOR}) - ->Threads(1) - ->Threads(2) - ->Threads(4) - ->Threads(8) - ->Threads(16) - ->Threads(24) - ->Unit(benchmark::kSecond); - */ + ->UseRealTime(); + /* + ->Args({96*4, arrow::Compression::FASTPFOR}) + ->Args({96*8, arrow::Compression::FASTPFOR}) + ->Args({96*16, arrow::Compression::FASTPFOR}) + ->Args({96*32, arrow::Compression::FASTPFOR}) + benchmark::RegisterBenchmark("BenchmarkShuffleSplit::IterateScan", bck) + ->Iterations(1) + ->Args({4096, arrow::Compression::FASTPFOR}) + ->Threads(1) + ->Threads(2) + ->Threads(4) + ->Threads(8) + ->Threads(16) + ->Threads(24) + ->Unit(benchmark::kSecond); + */ benchmark::Initialize(&argc, argv); benchmark::RunSpecifiedBenchmarks(); benchmark::Shutdown(); -} \ No newline at end of file +} diff --git a/native-sql-engine/cpp/src/shuffle/splitter.cc b/native-sql-engine/cpp/src/shuffle/splitter.cc index 3c7d8eeee..e4571cf36 100644 --- a/native-sql-engine/cpp/src/shuffle/splitter.cc +++ b/native-sql-engine/cpp/src/shuffle/splitter.cc @@ -26,6 +26,7 @@ #include #include #include +#include #include #include @@ -338,14 +339,19 @@ arrow::Status Splitter::AllocateBufferFromPool(std::shared_ptr& b // make size 64byte aligned auto reminder = size & 0x3f; size += (64 - reminder) & ((reminder == 0) - 1); - if (size > SPLIT_BUFFER_SIZE) { + + if (size > combine_buffer_size_) { ARROW_ASSIGN_OR_RAISE(buffer, arrow::AllocateResizableBuffer(size, options_.memory_pool)); return arrow::Status::OK(); } else if (combine_buffer_->capacity() - combine_buffer_->size() < size) { // memory pool is not enough - ARROW_ASSIGN_OR_RAISE(combine_buffer_, arrow::AllocateResizableBuffer( - SPLIT_BUFFER_SIZE, options_.memory_pool)); + ARROW_ASSIGN_OR_RAISE( + combine_buffer_, + arrow::AllocateResizableBuffer(combine_buffer_size_, options_.memory_pool)); + madvise(combine_buffer_->mutable_data(), combine_buffer_size_, /*MADV_HUGEPAGE */ 14); + madvise(combine_buffer_->mutable_data(), combine_buffer_size_, MADV_WILLNEED); + combine_buffer_->Resize(0, /*shrink_to_fit = */ false); } buffer = arrow::SliceMutableBuffer(combine_buffer_, combine_buffer_->size(), size); @@ -514,7 +520,7 @@ arrow::Status Splitter::CacheRecordBatch(int32_t partition_id, bool reset_buffer buffers[1]->data())[0] : reinterpret_cast( buffers[1]->data())[0]; - ARROW_CHECK_EQ(dst_offset0, 0); + ARROW_CHECK_EQ(dst_offset0, 0) << dst_offset0; if (reset_buffers) { partition_validity_addrs_[fixed_width_col_cnt_ + binary_idx][partition_id] = @@ -629,11 +635,10 @@ arrow::Status Splitter::AllocatePartitionBuffers(int32_t partition_id, int32_t n } std::shared_ptr offset_buffer; + std::shared_ptr value_buffer; std::shared_ptr validity_buffer = nullptr; auto value_buf_size = binary_array_empirical_size_[binary_idx] * new_size + 1024; - ARROW_ASSIGN_OR_RAISE( - std::shared_ptr value_buffer, - arrow::AllocateResizableBuffer(value_buf_size, options_.memory_pool)); + ARROW_RETURN_NOT_OK(AllocateBufferFromPool(value_buffer, value_buf_size)); ARROW_RETURN_NOT_OK( AllocateBufferFromPool(offset_buffer, new_size * sizeof_binary_offset + 1)); // set the first offset to 0 @@ -817,6 +822,7 @@ Splitter::row_offset_type Splitter::CalculateSplitBatchSize( reinterpret_cast( offsetbuf)[num_rows] - reinterpret_cast(offsetbuf)[0]; + size_per_row += sizeof(arrow::StringType::offset_type); break; case arrow::LargeBinaryType::type_id: case arrow::LargeStringType::type_id: @@ -824,6 +830,7 @@ Splitter::row_offset_type Splitter::CalculateSplitBatchSize( offsetbuf)[num_rows] - reinterpret_cast( offsetbuf)[0]; + size_per_row += sizeof(arrow::LargeStringType::offset_type); break; } binary_array_empirical_size_[i - fixed_width_col_cnt_] = @@ -833,21 +840,27 @@ Splitter::row_offset_type Splitter::CalculateSplitBatchSize( } } } - size_per_row = std::accumulate(binary_array_empirical_size_.begin(), - binary_array_empirical_size_.end(), 0); + size_per_row += std::accumulate(binary_array_empirical_size_.begin(), + binary_array_empirical_size_.end(), 0); - for (auto col = 0; col < array_idx_.size(); ++col) { + for (auto col = 0; col < fixed_width_col_cnt_; ++col) { auto col_idx = array_idx_[col]; - if (col_idx < fixed_width_col_cnt_) - size_per_row += arrow::bit_width(column_type_id_[col_idx]->id()) >> 3; + size_per_row += arrow::bit_width(column_type_id_[col_idx]->id()) >> 3; } + // add validity buffer + size_per_row += (array_idx_.size() >> 3) + 1; + int64_t prealloc_row_cnt = options_.offheap_per_task > 0 && size_per_row > 0 ? options_.offheap_per_task / size_per_row / num_partitions_ >> 2 : options_.buffer_size; prealloc_row_cnt = std::min(prealloc_row_cnt, (int64_t)options_.buffer_size); + combine_buffer_size_ = 1L * prealloc_row_cnt * size_per_row * num_partitions_; + // each binary buffer is allocated 1024 bytes more + combine_buffer_size_ += (array_idx_.size() - fixed_width_col_cnt_) * 1024; + return (row_offset_type)prealloc_row_cnt; } @@ -1232,21 +1245,42 @@ arrow::Status Splitter::SplitBinaryType(const uint8_t* src_addr, const T* src_of value_offset = dst_offset_base[x + 1] = value_offset + strlength; if (ARROW_PREDICT_FALSE(value_offset >= capacity)) { // allocate value buffer again - // enlarge the buffer by 1.5x - capacity = capacity + std::max((capacity >> 1), (uint64_t)strlength); + // calculated the buffer size based on already filled buffer + uint64_t size_per_row = value_offset / (partition_buffer_idx_base_[pid] + x) + 1; + combine_buffer_size_ = combine_buffer_size_ + + (size_per_row - binary_array_empirical_size_[binary_idx]) * + partition_buffer_size_[pid]; + binary_array_empirical_size_[binary_idx] = + std::max(binary_array_empirical_size_[binary_idx], size_per_row); + uint64_t new_size = + binary_array_empirical_size_[binary_idx] * partition_buffer_size_[pid] + 1024; + capacity = std::max(new_size, capacity + strlength); auto value_buffer = std::static_pointer_cast( partition_buffers_[fixed_width_col_cnt_ + binary_idx][pid][2]); - value_buffer->Reserve(capacity); + // if it's allocated from combine buffer, we need to allocate a new buffer and + // copy the data there + if (value_buffer->parent() != NULL) { + ARROW_ASSIGN_OR_RAISE(value_buffer, arrow::AllocateResizableBuffer( + capacity, options_.memory_pool)); + memcpy(value_buffer->mutable_data(), dst_addrs[pid].valueptr, + dst_addrs[pid].value_capacity); + partition_buffers_[fixed_width_col_cnt_ + binary_idx][pid][2] = value_buffer; + } else { + // if it's allocated from memory pool directly, reserve more memory + value_buffer->Reserve(capacity); + } - dst_addrs[pid].valueptr = value_buffer->mutable_data(); - dst_addrs[pid].value_capacity = capacity; - dst_value_base = dst_addrs[pid].valueptr + value_offset - strlength; std::cout << " value buffer resized colid = " << binary_idx << " dst_start " << dst_offset_base[x] << " dst_end " << dst_offset_base[x + 1] - << " old size = " << capacity << " new size = " << capacity - << " row = " << partition_buffer_idx_base_[pid] + << " old size = " << dst_addrs[pid].value_capacity + << " new size = " << capacity + << " row = " << (partition_buffer_idx_base_[pid] + x) << " strlen = " << strlength << std::endl; + + dst_addrs[pid].valueptr = value_buffer->mutable_data(); + dst_addrs[pid].value_capacity = capacity; + dst_value_base = dst_addrs[pid].valueptr + value_offset - strlength; } auto value_src_ptr = src_addr + src_offset_addr[src_offset]; //#if NATIVE_AVX512 == ON diff --git a/native-sql-engine/cpp/src/shuffle/splitter.h b/native-sql-engine/cpp/src/shuffle/splitter.h index 6cd64fd23..0f111b134 100644 --- a/native-sql-engine/cpp/src/shuffle/splitter.h +++ b/native-sql-engine/cpp/src/shuffle/splitter.h @@ -228,6 +228,7 @@ class Splitter { // slice the buffer for each reducer's column, in this way we can combine into large // page std::shared_ptr combine_buffer_; + uint64_t combine_buffer_size_; // partid std::vector>>