[VL] Support celeborn sort based shuffle #5675

kerwin-zk · 2024-05-09T12:33:57Z

What changes were proposed in this pull request?

#4048

How was this patch tested?

(Please explain how this patch was tested. E.g. unit tests, integration tests, manual tests)

(If this patch involves UI changes, please attach a screenshot; otherwise, remove this)

github-actions · 2024-05-09T12:34:15Z

Thanks for opening a pull request!

Could you open an issue for this pull request on Github Issues?

https://github.com/apache/incubator-gluten/issues

Then could you also rename commit message and pull request title in the following format?

[GLUTEN-${ISSUES_ID}][COMPONENT]feat/fix: ${detailed message}

See also:

Other pull requests

github-actions · 2024-05-09T12:34:29Z

Run Gluten Clickhouse CI

github-actions · 2024-05-09T12:34:46Z

Run Gluten Clickhouse CI

github-actions · 2024-05-09T12:46:03Z

Run Gluten Clickhouse CI

github-actions · 2024-05-09T15:30:47Z

Run Gluten Clickhouse CI

github-actions · 2024-05-09T16:45:20Z

Run Gluten Clickhouse CI

FelixYBW · 2024-05-09T19:16:19Z

Thank you for the PR! Waiting for long time.

FelixYBW · 2024-05-09T19:19:34Z

cpp/velox/shuffle/VeloxShuffleWriter.h

@@ -132,6 +132,8 @@ class VeloxShuffleWriter final : public ShuffleWriter {

  arrow::Status split(std::shared_ptr<ColumnarBatch> cb, int64_t memLimit) override;

+  arrow::Status sort(std::shared_ptr<ColumnarBatch> cb, int64_t memLimit) override;


We may separate the sort based and hash based shuffle in different files

@FelixYBW After Sort's proposal is reviewed, I'll refactor the code structure to meet the requirements of the different files.

@kerwin-zk After this patch, part of the existing APIs are only used for hash-based shuffle, and the newly added ones are only used for sort-based shuffle, which leads to a mixing of the code. It would be nice to separate the APIs for both writer & reader.

@marin-ma Shuffle Writer's api has been split into sort and hash, for Shuffle Reader, it just adds Velox Stream parsing, and doesn't differentiate between sort and hash, the core logic is the same, I don't feel the need to split it out separately.

github-actions · 2024-05-10T01:01:08Z

Run Gluten Clickhouse CI

github-actions · 2024-05-10T01:02:52Z

Run Gluten Clickhouse CI

zhouyuan · 2024-05-10T04:09:58Z

cpp/velox/shuffle/VeloxShuffleWriter.cc

-arrow::Status VeloxShuffleWriter::stop() {
+arrow::Status VeloxShuffleWriter::doSort(facebook::velox::RowVectorPtr rv, int64_t memLimit) {
+  currentInputColumnBytes_ += rv->estimateFlatSize();
+  batches_.push_back(rv);


QQ: looks like this only enqueue the batches, the sort is done in Celeborn side?

@zhouyuan In the HashPartitioner's compute, it will encapsulate the information of each row belonging to the same PartitionId in the RowVector into an int, and put it into the Vector of the map, which is equivalent to replacing the Sort with a HashMap.

github-actions · 2024-05-10T08:46:58Z

Run Gluten Clickhouse CI

github-actions · 2024-05-10T08:52:47Z

Run Gluten Clickhouse CI

github-actions · 2024-05-10T11:34:44Z

Run Gluten Clickhouse CI

github-actions · 2024-05-10T12:01:02Z

Run Gluten Clickhouse CI

marin-ma

Could you add cpp UT for sort-based shuffle? If it's possible to implement another LocalRssClient for sort-based shuffle and use it in native tests, we can reuse the existing tests by adding a new test param in VeloxShuffleWriterTest

marin-ma · 2024-05-13T06:03:36Z

cpp/core/jni/JniWrapper.cc

@@ -148,6 +149,69 @@ class JavaInputStreamAdaptor final : public arrow::io::InputStream {
  bool closed_ = false;
 };

+class JavaInputStreamVeloxWrapper final : public JavaInputStreamWrapper {


Can we use JavaInputStreamAdapter?

The existing JavaInputStreamAdapter is defined here

incubator-gluten/cpp/core/jni/JniWrapper.cc

Line 79 in 182a029

class JavaInputStreamAdaptor final : public arrow::io::InputStream {

Can we use this class?

@marin-ma done

marin-ma · 2024-05-13T06:19:32Z

cpp/velox/compute/VeloxRuntime.cc

-  return std::make_shared<VeloxShuffleReader>(std::move(deserializerFactory));
+      schema,
+      std::move(codec),
+      options.compressionTypeStr,


As it's in Velox backend, we can put the str->kind conversion here annd store the velox::common::CompressionKind instead of the compression string.

github-actions · 2024-05-13T13:00:31Z

Run Gluten Clickhouse CI

github-actions · 2024-05-13T16:06:05Z

Run Gluten Clickhouse CI

github-actions · 2024-05-14T02:07:54Z

Run Gluten Clickhouse CI

github-actions · 2024-05-14T02:41:25Z

Run Gluten Clickhouse CI

github-actions · 2024-05-14T08:23:24Z

Run Gluten Clickhouse CI

github-actions · 2024-05-14T08:24:37Z

Run Gluten Clickhouse CI

github-actions · 2024-05-14T08:36:22Z

Run Gluten Clickhouse CI

marin-ma

A few comments.

marin-ma · 2024-05-14T05:39:59Z

cpp/core/shuffle/HashPartitioner.cc

+    std::unordered_map<int32_t, std::vector<int64_t>>& rowVectorIndexMap) {
+  auto index = static_cast<int64_t>(vectorIndex) << 32;
+  for (auto i = 0; i < numRows; ++i) {
+    auto pid = pidArr[i] % numPartitions_;


These are the same in the function above. Could you extract the common part?

marin-ma · 2024-05-14T09:06:44Z

cpp/velox/shuffle/VeloxSortBasedShuffleWriter.cc

+    facebook::velox::OStreamOutputStream* out,
+    facebook::velox::RowTypePtr* rowTypePtr) {
+  int64_t rawSize = batch_->size();
+  batch_->flush(out);


According to PrestoSerializer implementation, looks like there's an unavoidable extra memory copy here, from the intermediate buffer used by batch_->append to std::ostringstream* output. And the memory used by std::ostringstream* output is not being tracked. Is there any way to track the memory usage?

We can use IOBufOutputStream so that the memory can be tracked by velox memory pool.

@kerwin-zk Just found that the data stored in IOBufOutputStream might not be contiguous. Let's keep the current implementation as for now. I will find a better way to resolve it.

@marin-ma That's the problem. If it's not contiguous, there's an extra copy.

marin-ma · 2024-05-14T09:12:12Z

cpp/velox/shuffle/VeloxSortBasedShuffleWriter.cc

+      for (int start = 0; start < outputSize; start++) {
+        const int32_t vectorIndex = static_cast<int32_t>(rowVectorIndex.at(start) >> 32);
+        const int32_t rowIndex = static_cast<int32_t>(rowVectorIndex.at(start) & 0xFFFFFFFFLL);
+        groupedIndices[vectorIndex].push_back({rowIndex, 1});


The mapping from the input vector index + row index is already ordered. Can we directly build the index range in Partitioner::compute? Then the adjacent row index of the same partition id can be merged.

@marin-ma Partitioner::compute belongs to the core part of the cpp, it is not appropriate to introduce Velox-related IndexRange, I will implement this part of the logic in evictRowVector.

github-actions · 2024-05-16T06:11:10Z

Run Gluten Clickhouse CI

github-actions · 2024-05-16T09:48:41Z

Run Gluten Clickhouse CI

marin-ma

A few comments. Thanks.

marin-ma · 2024-05-16T09:09:49Z

cpp/core/shuffle/ShuffleReader.cc

@@ -16,6 +16,9 @@
 */

 #include "ShuffleReader.h"
+
+#include <jni/JniCommon.h>


Unintended change?

marin-ma · 2024-05-16T09:12:03Z

cpp/velox/tests/VeloxShuffleWriterTest.cc

-      params.push_back(ShuffleTestParams{PartitionWriterType::kRss, compression, compressionThreshold, 0});
+      params.push_back(ShuffleTestParams{
+          ShuffleWriterType::kHashShuffle, PartitionWriterType::kRss, compression, compressionThreshold, 0});
+      params.push_back(ShuffleTestParams{


Looks like the "compressionThreshold" is not used in the "sort + rss" code path. We can move this combination out of this loop to reduce the number of test cases.

marin-ma · 2024-05-16T09:14:33Z

cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h

+    std::shared_ptr<VeloxShuffleWriter> shuffleWriter;
+    if (GetParam().shuffleWriterType == kHashShuffle) {
+      GLUTEN_ASSIGN_OR_THROW(
+          shuffleWriter,
+          VeloxHashBasedShuffleWriter::create(
+              kNumPartitions, std::move(partitionWriter), std::move(shuffleWriterOptions_), pool_, arrowPool));
+    } else if (
+        GetParam().shuffleWriterType == kSortShuffle && GetParam().partitionWriterType == PartitionWriterType::kRss) {
+      GLUTEN_ASSIGN_OR_THROW(
+          shuffleWriter,
+          VeloxSortBasedShuffleWriter::create(
+              kNumPartitions, std::move(partitionWriter), std::move(shuffleWriterOptions_), pool_, arrowPool));
+    }


Please extract this logic into a function to avoid copy-paste.

marin-ma · 2024-05-16T09:17:04Z

gluten-data/src/main/java/org/apache/gluten/vectorized/ShuffleWriterJniWrapper.java

      Object pusher,
-      String partitionWriterType);
+      String partitionWriterType,
+      String shuffleWriterType);


Maybe we can consider separating the JNI call for hash/sort shuffle in the future.

marin-ma · 2024-05-16T10:46:23Z

@kerwin-zk It's weird that I cannot reproduce the UT failure on my dev environment (ubuntu 22.04), but can reproduce that in the docker container (centos stream 8).

kerwin-zk · 2024-05-16T11:10:53Z

@kerwin-zk It's weird that I cannot reproduce the UT failure on my dev environment (ubuntu 22.04), but can reproduce that in the docker container (centos stream 8).

@marin-ma Yes, I am checking this issue today. I tested it using centos8 and ubuntu2004, but the local development environment could not reproduce it.

marin-ma · 2024-05-16T11:15:03Z

I tested it using centos8 and ubuntu2004, but the local development environment could not reproduce it.

@kerwin-zk Do you mean it cannot be reproduced on centos8? What GCC version do you use? It's gcc-9 in docker https://github.com/apache/incubator-gluten/blob/main/.github/workflows/velox_docker.yml#L524

kerwin-zk · 2024-05-16T11:23:53Z

I tested it using centos8 and ubuntu2004, but the local development environment could not reproduce it.

@kerwin-zk Do you mean it cannot be reproduced on centos8? What GCC version do you use? It's gcc-9 in docker https://github.com/apache/incubator-gluten/blob/main/.github/workflows/velox_docker.yml#L524

@marin-ma centos8 gcc10

github-actions · 2024-05-16T16:07:03Z

Run Gluten Clickhouse CI

github-actions · 2024-05-16T16:09:21Z

Run Gluten Clickhouse CI

github-actions · 2024-05-16T18:37:15Z

Run Gluten Clickhouse CI

github-actions · 2024-05-17T01:07:17Z

Run Gluten Clickhouse CI

github-actions · 2024-05-17T02:30:30Z

Run Gluten Clickhouse CI

marin-ma

LGTM. Some minor comments. Thanks!

marin-ma · 2024-05-17T02:36:41Z

cpp/core/shuffle/HashPartitioner.cc

@@ -16,9 +16,28 @@
 */

 #include "shuffle/HashPartitioner.h"
+#include <iostream>


Please remove.

marin-ma · 2024-05-17T02:38:38Z

cpp/core/shuffle/LocalPartitionWriter.cc

@@ -541,6 +541,10 @@ arrow::Status LocalPartitionWriter::evict(
  return arrow::Status::OK();
 }

+arrow::Status LocalPartitionWriter::evict(uint32_t partitionId, int64_t rawSize, const char* data, int64_t length) {
+  return arrow::Status::OK();


Return arrow::Status::NotImplemented("Invalid code path for local shuffle writer: sort based is not supported.");

marin-ma · 2024-05-17T02:41:00Z

cpp/core/shuffle/Options.h

 enum PartitionWriterType { kLocal, kRss };

 struct ShuffleReaderOptions {
  arrow::Compression::type compressionType = arrow::Compression::type::LZ4_FRAME;
+  std::string compressionTypeStr = "lz4";


Please add static const std::string kDefaultCompressionTypeStr = "lz4"; in above code and use kDefaultCompressionTypeStr here.

marin-ma · 2024-05-17T02:41:38Z

cpp/core/shuffle/Options.h

@@ -56,6 +63,7 @@ struct PartitionWriterOptions {

  int32_t compressionThreshold = kDefaultCompressionThreshold;
  arrow::Compression::type compressionType = arrow::Compression::LZ4_FRAME;
+  std::string compressionTypeStr = "lz4";


kDefaultCompressionTypeStr

marin-ma · 2024-05-17T02:43:26Z

cpp/velox/memory/VeloxColumnarBatch.h

@@ -46,7 +46,7 @@ class VeloxColumnarBatch final : public ColumnarBatch {
  facebook::velox::RowVectorPtr getRowVector() const;
  facebook::velox::RowVectorPtr getFlattenedRowVector();

- private:
+ public:


Unintended change?

marin-ma · 2024-05-17T02:45:23Z

cpp/velox/shuffle/VeloxSortBasedShuffleWriter.h

+#if VELOX_SHUFFLE_WRITER_PRINT
+
+#define VsPrint Print
+#define VsPrintLF PrintLF
+#define VsPrintSplit PrintSplit
+#define VsPrintSplitLF PrintSplitLF
+#define VsPrintVectorRange PrintVectorRange
+#define VS_PRINT PRINT
+#define VS_PRINTLF PRINTLF
+#define VS_PRINT_FUNCTION_NAME PRINT_FUNCTION_NAME
+#define VS_PRINT_FUNCTION_SPLIT_LINE PRINT_FUNCTION_SPLIT_LINE
+#define VS_PRINT_CONTAINER PRINT_CONTAINER
+#define VS_PRINT_CONTAINER_TO_STRING PRINT_CONTAINER_TO_STRING
+#define VS_PRINT_CONTAINER_2_STRING PRINT_CONTAINER_2_STRING
+#define VS_PRINT_VECTOR_TO_STRING PRINT_VECTOR_TO_STRING
+#define VS_PRINT_VECTOR_2_STRING PRINT_VECTOR_2_STRING
+#define VS_PRINT_VECTOR_MAPPING PRINT_VECTOR_MAPPING
+
+#else // VELOX_SHUFFLE_WRITER_PRINT
+
+#define VsPrint(...) // NOLINT
+#define VsPrintLF(...) // NOLINT
+#define VsPrintSplit(...) // NOLINT
+#define VsPrintSplitLF(...) // NOLINT
+#define VsPrintVectorRange(...) // NOLINT
+#define VS_PRINT(a)
+#define VS_PRINTLF(a)
+#define VS_PRINT_FUNCTION_NAME()
+#define VS_PRINT_FUNCTION_SPLIT_LINE()
+#define VS_PRINT_CONTAINER(c)
+#define VS_PRINT_CONTAINER_TO_STRING(c)
+#define VS_PRINT_CONTAINER_2_STRING(c)
+#define VS_PRINT_VECTOR_TO_STRING(v)
+#define VS_PRINT_VECTOR_2_STRING(v)
+#define VS_PRINT_VECTOR_MAPPING(v)
+
+#endif // end of VELOX_SHUFFLE_WRITER_PRINT


Are these macros used by the sort shuffle writer? If not, please remove.

github-actions · 2024-05-17T02:58:00Z

Run Gluten Clickhouse CI

GlutenPerfBot · 2024-05-17T07:41:17Z

===== Performance report for TPCH SF2000 with Velox backend, for reference only ====

query	log/native_5675_time.csv	log/native_master_05_13_2024_33f993554_time.csv	difference	percentage
q1	35.86	36.42	0.562	101.57%
q2	24.03	23.70	-0.333	98.62%
q3	36.79	36.61	-0.177	99.52%
q4	36.84	37.83	0.983	102.67%
q5	68.64	70.62	1.980	102.88%
q6	7.70	7.47	-0.232	96.98%
q7	84.61	85.08	0.465	100.55%
q8	85.04	84.30	-0.735	99.14%
q9	123.72	121.00	-2.715	97.81%
q10	44.98	45.37	0.383	100.85%
q11	19.81	19.49	-0.317	98.40%
q12	25.88	26.52	0.644	102.49%
q13	56.41	55.08	-1.331	97.64%
q14	17.54	21.53	3.985	122.71%
q15	31.38	29.02	-2.361	92.47%
q16	14.17	14.00	-0.167	98.82%
q17	103.78	103.62	-0.160	99.85%
q18	145.89	146.24	0.342	100.23%
q19	13.42	13.35	-0.067	99.50%
q20	28.57	29.06	0.490	101.72%
q21	279.91	282.54	2.632	100.94%
q22	15.84	14.61	-1.227	92.26%
total	1300.80	1303.44	2.643	100.20%

kerwin-zk force-pushed the wip-celeborn-sort branch from a34f4f9 to 3bed73f Compare May 9, 2024 12:34

kerwin-zk force-pushed the wip-celeborn-sort branch from 3bed73f to dbd0f80 Compare May 9, 2024 12:45

kerwin-zk force-pushed the wip-celeborn-sort branch from dbd0f80 to d6cf16d Compare May 9, 2024 15:30

kerwin-zk force-pushed the wip-celeborn-sort branch from d6cf16d to b140199 Compare May 9, 2024 16:44

FelixYBW reviewed May 9, 2024

View reviewed changes

FelixYBW requested a review from marin-ma May 9, 2024 19:20

kerwin-zk force-pushed the wip-celeborn-sort branch from 828b059 to 81771a4 Compare May 10, 2024 01:02

zhouyuan reviewed May 10, 2024

View reviewed changes

marin-ma reviewed May 13, 2024

View reviewed changes

marin-ma reviewed May 14, 2024

View reviewed changes

kerwin-zk added 2 commits May 16, 2024 14:10

optimize index range

e6d44e6

[WIP][VL] Support celeborn sort based shuffle

f9d73e6

kerwin-zk force-pushed the wip-celeborn-sort branch from 2923db6 to f9d73e6 Compare May 16, 2024 06:10

[WIP][VL] Support celeborn sort based shuffle

71c6c1b

marin-ma reviewed May 16, 2024

View reviewed changes

[WIP][VL] Support celeborn sort based shuffle

7c13026

[WIP][VL] Support celeborn sort based shuffle

89dde0b

debug

b1cf3c6

debug

1ca47e1

[WIP][VL] Support celeborn sort based shuffle

9cf13c3

marin-ma reviewed May 17, 2024

View reviewed changes

marin-ma changed the title ~~[WIP][VL] Support celeborn sort based shuffle~~ [VL] Support celeborn sort based shuffle May 17, 2024

[WIP][VL] Support celeborn sort based shuffle

4c4298a

This was referenced May 17, 2024

[VL] Add BufferedOutputStream to track the memory usage in PrestoSerializer #5785

Merged

[VL] Refine evict logic in sort shuffle writer #5786

Merged

marin-ma merged commit dfac04f into apache:main May 17, 2024
43 checks passed

xumingming mentioned this pull request May 17, 2024

[CORE] Build failed for macOS #5792

Closed

		@@ -132,6 +132,8 @@ class VeloxShuffleWriter final : public ShuffleWriter {

		arrow::Status split(std::shared_ptr<ColumnarBatch> cb, int64_t memLimit) override;

		arrow::Status sort(std::shared_ptr<ColumnarBatch> cb, int64_t memLimit) override;

[VL] Support celeborn sort based shuffle #5675

[VL] Support celeborn sort based shuffle #5675

Conversation

kerwin-zk commented May 9, 2024

What changes were proposed in this pull request?

How was this patch tested?

github-actions bot commented May 9, 2024

github-actions bot commented May 9, 2024

github-actions bot commented May 9, 2024

github-actions bot commented May 9, 2024

github-actions bot commented May 9, 2024

github-actions bot commented May 9, 2024

FelixYBW commented May 9, 2024

FelixYBW May 9, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented May 10, 2024

github-actions bot commented May 10, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented May 10, 2024

github-actions bot commented May 10, 2024

github-actions bot commented May 10, 2024

github-actions bot commented May 10, 2024

marin-ma left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented May 13, 2024

github-actions bot commented May 13, 2024

github-actions bot commented May 14, 2024

github-actions bot commented May 14, 2024

github-actions bot commented May 14, 2024

github-actions bot commented May 14, 2024

github-actions bot commented May 14, 2024

marin-ma left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

marin-ma May 14, 2024 • edited

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented May 16, 2024

github-actions bot commented May 16, 2024

marin-ma left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

marin-ma commented May 16, 2024 • edited

kerwin-zk commented May 16, 2024

marin-ma commented May 16, 2024 • edited

kerwin-zk commented May 16, 2024

github-actions bot commented May 16, 2024

github-actions bot commented May 16, 2024

github-actions bot commented May 16, 2024

github-actions bot commented May 17, 2024

github-actions bot commented May 17, 2024

marin-ma left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

github-actions bot commented May 17, 2024

GlutenPerfBot commented May 17, 2024

FelixYBW May 9, 2024 •

edited

marin-ma May 14, 2024 •

edited

marin-ma commented May 16, 2024 •

edited

marin-ma commented May 16, 2024 •

edited