[VL][CI] Fix fallback to columnar shuffle in celeborn test (#4474)

apache · Jan 29, 2024 · 7ed8e37 · 7ed8e37
1 parent 242cca6
commit 7ed8e37
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 17 deletions.
diff --git a/.github/workflows/velox_be.yml b/.github/workflows/velox_be.yml
@@ -372,10 +372,12 @@ jobs:
       - name: TPC-H SF1.0 && TPC-DS SF10.0 Parquet local spark3.2 with Celeborn
         run: |
           $PATH_TO_GLUTEN_TE/$OS_IMAGE_NAME/gha/gha-checkout/exec.sh \
-          'mv /opt/apache-celeborn-0.3.0-incubating-bin/conf/celeborn-env.sh.template /opt/apache-celeborn-0.3.0-incubating-bin/conf/celeborn-env.sh && \
-          echo -e "CELEBORN_MASTER_MEMORY=4g\nCELEBORN_WORKER_MEMORY=4g\nCELEBORN_WORKER_OFFHEAP_MEMORY=8g" > /opt/apache-celeborn-0.3.0-incubating-bin/conf/celeborn-env.sh && \
-          echo -e "celeborn.worker.commitFiles.threads 128\nceleborn.worker.sortPartition.threads 64" > /opt/apache-celeborn-0.3.0-incubating-bin/conf/celeborn-defaults.conf \
-          && bash /opt/apache-celeborn-0.3.0-incubating-bin/sbin/start-master.sh && bash /opt/apache-celeborn-0.3.0-incubating-bin/sbin/start-worker.sh && \
+          'wget https://dlcdn.apache.org/incubator/celeborn/celeborn-0.3.0-incubating/apache-celeborn-0.3.0-incubating-bin.tgz && \
+          tar xzf apache-celeborn-0.3.0-incubating-bin.tgz && cd apache-celeborn-0.3.0-incubating-bin && \
+          mv ./conf/celeborn-env.sh.template ./conf/celeborn-env.sh && \
+          echo -e "CELEBORN_MASTER_MEMORY=4g\nCELEBORN_WORKER_MEMORY=4g\nCELEBORN_WORKER_OFFHEAP_MEMORY=8g" > ./conf/celeborn-env.sh && \
+          echo -e "celeborn.worker.commitFiles.threads 128\nceleborn.worker.sortPartition.threads 64" > ./conf/celeborn-defaults.conf \
+          && bash ./sbin/start-master.sh && bash ./sbin/start-worker.sh && \
           cd /opt/gluten/tools/gluten-it && mvn clean install -Pspark-3.2,rss \
           && GLUTEN_IT_JVM_ARGS=-Xmx5G sbin/gluten-it.sh queries-compare \
             --local --preset=velox-with-celeborn --benchmark-type=h --error-on-memleak --off-heap-size=10g -s=1.0 --threads=16 --iterations=1 \

diff --git a/cpp/velox/memory/VeloxColumnarBatch.cc b/cpp/velox/memory/VeloxColumnarBatch.cc
@@ -16,7 +16,6 @@
  */
 #include "VeloxColumnarBatch.h"
 #include "compute/VeloxRuntime.h"
-#include "utils/Timer.h"
 #include "utils/VeloxArrowUtils.h"
 #include "velox/row/UnsafeRowFast.h"
 #include "velox/type/Type.h"
@@ -48,17 +47,20 @@ void VeloxColumnarBatch::ensureFlattened() {
   if (flattened_ != nullptr) {
     return;
   }
-
-  ScopedTimer timer(&exportNanos_);
+  auto startTime = std::chrono::steady_clock::now();
+  // Make sure to load lazy vector if not loaded already.
   for (auto& child : rowVector_->children()) {
-    facebook::velox::BaseVector::flattenVector(child);
-    if (child->isLazy()) {
-      child = child->as<facebook::velox::LazyVector>()->loadedVectorShared();
-      VELOX_DCHECK_NOT_NULL(child);
-    }
+    child->loadedVector();
   }
-  // Invalid the original rowVector_ after being flattened.
-  flattened_ = std::move(rowVector_);
+
+  // Perform copy to flatten dictionary vectors.
+  velox::RowVectorPtr copy = std::dynamic_pointer_cast<velox::RowVector>(
+      velox::BaseVector::create(rowVector_->type(), rowVector_->size(), rowVector_->pool()));
+  copy->copy(rowVector_.get(), 0, 0, rowVector_->size());
+  flattened_ = copy;
+  auto endTime = std::chrono::steady_clock::now();
+  auto duration = std::chrono::duration_cast<std::chrono::nanoseconds>(endTime - startTime).count();
+  exportNanos_ += duration;
 }
 
 std::shared_ptr<ArrowSchema> VeloxColumnarBatch::exportArrowSchema() {
@@ -81,7 +83,6 @@ int64_t VeloxColumnarBatch::numBytes() {
 }
 
 velox::RowVectorPtr VeloxColumnarBatch::getRowVector() const {
-  VELOX_CHECK_NOT_NULL(rowVector_);
   return rowVector_;
 }
 

diff --git a/tools/gluten-te/ubuntu/dockerfile-buildenv b/tools/gluten-te/ubuntu/dockerfile-buildenv
@@ -115,8 +115,6 @@ RUN cd /opt/spark331 && ./build/mvn -Pyarn -DskipTests clean install
 RUN cd /opt && git clone --depth 1 --branch v3.4.2 https://github.com/apache/spark.git spark342
 RUN cd /opt/spark342 && ./build/mvn -Pyarn -DskipTests clean install
 
-RUN cd /opt && git clone --depth 1 --branch v0.3.0-incubating https://github.com/apache/incubator-celeborn.git apache-celeborn-0.3.0-incubating-bin
-
 # Prepare entry command
 COPY scripts/cmd.sh /root/.cmd.sh
 CMD ["/root/.cmd.sh"]