This is an automated email from the ASF dual-hosted git repository.
marong pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/incubator-gluten.git
The following commit(s) were added to refs/heads/main by this push:
new 7af129cba6 [GLUTEN-8481][VL] Clean up shuffle reader cpp code
7af129cba6 is described below
commit 7af129cba65f3987fa899f83d7560c19014a6943
Author: Rong Ma <[email protected]>
AuthorDate: Fri Jan 10 10:04:53 2025 +0800
[GLUTEN-8481][VL] Clean up shuffle reader cpp code
---
cpp/core/CMakeLists.txt | 1 -
cpp/core/jni/JniWrapper.cc | 1 -
cpp/core/shuffle/ShuffleReader.cc | 55 ----------------------
cpp/core/shuffle/ShuffleReader.h | 49 ++-----------------
cpp/velox/compute/VeloxRuntime.cc | 2 +-
cpp/velox/shuffle/VeloxShuffleReader.cc | 40 ++++++++++------
cpp/velox/shuffle/VeloxShuffleReader.h | 33 ++++++++-----
cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h | 2 +-
8 files changed, 52 insertions(+), 131 deletions(-)
diff --git a/cpp/core/CMakeLists.txt b/cpp/core/CMakeLists.txt
index 5a5eeac354..d52ad74b99 100644
--- a/cpp/core/CMakeLists.txt
+++ b/cpp/core/CMakeLists.txt
@@ -120,7 +120,6 @@ set(SPARK_COLUMNAR_PLUGIN_SRCS
shuffle/RandomPartitioner.cc
shuffle/RoundRobinPartitioner.cc
shuffle/ShuffleMemoryPool.cc
- shuffle/ShuffleReader.cc
shuffle/ShuffleWriter.cc
shuffle/SinglePartitioner.cc
shuffle/Spill.cc
diff --git a/cpp/core/jni/JniWrapper.cc b/cpp/core/jni/JniWrapper.cc
index b4359bab41..4e1d1f09fd 100644
--- a/cpp/core/jni/JniWrapper.cc
+++ b/cpp/core/jni/JniWrapper.cc
@@ -1132,7 +1132,6 @@ JNIEXPORT void JNICALL
Java_org_apache_gluten_vectorized_ShuffleReaderJniWrapper
jlong shuffleReaderHandle) {
JNI_METHOD_START
auto reader = ObjectStore::retrieve<ShuffleReader>(shuffleReaderHandle);
- GLUTEN_THROW_NOT_OK(reader->close());
ObjectStore::release(shuffleReaderHandle);
JNI_METHOD_END()
}
diff --git a/cpp/core/shuffle/ShuffleReader.cc
b/cpp/core/shuffle/ShuffleReader.cc
deleted file mode 100644
index ced80b3de1..0000000000
--- a/cpp/core/shuffle/ShuffleReader.cc
+++ /dev/null
@@ -1,55 +0,0 @@
-/*
- * Licensed to the Apache Software Foundation (ASF) under one or more
- * contributor license agreements. See the NOTICE file distributed with
- * this work for additional information regarding copyright ownership.
- * The ASF licenses this file to You under the Apache License, Version 2.0
- * (the "License"); you may not use this file except in compliance with
- * the License. You may obtain a copy of the License at
- *
- * http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-#include "ShuffleReader.h"
-#include "arrow/ipc/reader.h"
-#include "arrow/record_batch.h"
-#include "utils/Macros.h"
-
-#include <utility>
-
-#include "ShuffleSchema.h"
-
-namespace gluten {
-
-ShuffleReader::ShuffleReader(std::unique_ptr<DeserializerFactory> factory) :
factory_(std::move(factory)) {}
-
-std::shared_ptr<ResultIterator>
ShuffleReader::readStream(std::shared_ptr<arrow::io::InputStream> in) {
- return std::make_shared<ResultIterator>(factory_->createDeserializer(in));
-}
-
-arrow::Status ShuffleReader::close() {
- return arrow::Status::OK();
-}
-
-arrow::MemoryPool* ShuffleReader::getPool() const {
- return factory_->getPool();
-}
-
-int64_t ShuffleReader::getDecompressTime() const {
- return factory_->getDecompressTime();
-}
-
-ShuffleWriterType ShuffleReader::getShuffleWriterType() const {
- return factory_->getShuffleWriterType();
-}
-
-int64_t ShuffleReader::getDeserializeTime() const {
- return factory_->getDeserializeTime();
-}
-
-} // namespace gluten
diff --git a/cpp/core/shuffle/ShuffleReader.h b/cpp/core/shuffle/ShuffleReader.h
index 0f985c7da9..6e2b079fc7 100644
--- a/cpp/core/shuffle/ShuffleReader.h
+++ b/cpp/core/shuffle/ShuffleReader.h
@@ -17,63 +17,22 @@
#pragma once
-#include "memory/ColumnarBatch.h"
-
-#include <arrow/ipc/message.h>
-#include <arrow/ipc/options.h>
-
-#include "Options.h"
#include "compute/ResultIterator.h"
-#include "utils/Compression.h"
namespace gluten {
-class DeserializerFactory {
- public:
- virtual ~DeserializerFactory() = default;
-
- virtual std::unique_ptr<ColumnarBatchIterator>
createDeserializer(std::shared_ptr<arrow::io::InputStream> in) = 0;
-
- virtual arrow::MemoryPool* getPool() = 0;
-
- virtual int64_t getDecompressTime() = 0;
-
- virtual int64_t getDeserializeTime() = 0;
-
- virtual ShuffleWriterType getShuffleWriterType() = 0;
-};
-
class ShuffleReader {
public:
- explicit ShuffleReader(std::unique_ptr<DeserializerFactory> factory);
-
virtual ~ShuffleReader() = default;
// FIXME iterator should be unique_ptr or un-copyable singleton
- virtual std::shared_ptr<ResultIterator>
readStream(std::shared_ptr<arrow::io::InputStream> in);
-
- arrow::Status close();
-
- int64_t getDecompressTime() const;
-
- int64_t getIpcTime() const;
-
- int64_t getDeserializeTime() const;
-
- arrow::MemoryPool* getPool() const;
-
- ShuffleWriterType getShuffleWriterType() const;
+ virtual std::shared_ptr<ResultIterator>
readStream(std::shared_ptr<arrow::io::InputStream> in) = 0;
- protected:
- arrow::MemoryPool* pool_;
- int64_t decompressTime_ = 0;
- int64_t deserializeTime_ = 0;
+ virtual int64_t getDecompressTime() const = 0;
- ShuffleWriterType shuffleWriterType_;
+ virtual int64_t getDeserializeTime() const = 0;
- private:
- std::shared_ptr<arrow::Schema> schema_;
- std::unique_ptr<DeserializerFactory> factory_;
+ virtual arrow::MemoryPool* getPool() const = 0;
};
} // namespace gluten
diff --git a/cpp/velox/compute/VeloxRuntime.cc
b/cpp/velox/compute/VeloxRuntime.cc
index 20c3dec939..2a2ea929c1 100644
--- a/cpp/velox/compute/VeloxRuntime.cc
+++ b/cpp/velox/compute/VeloxRuntime.cc
@@ -251,7 +251,7 @@ std::shared_ptr<ShuffleReader>
VeloxRuntime::createShuffleReader(
auto codec = gluten::createArrowIpcCodec(options.compressionType,
options.codecBackend);
auto ctxVeloxPool = memoryManager()->getLeafMemoryPool();
auto veloxCompressionType =
facebook::velox::common::stringToCompressionKind(options.compressionTypeStr);
- auto deserializerFactory =
std::make_unique<gluten::VeloxColumnarBatchDeserializerFactory>(
+ auto deserializerFactory =
std::make_unique<gluten::VeloxShuffleReaderDeserializerFactory>(
schema,
std::move(codec),
veloxCompressionType,
diff --git a/cpp/velox/shuffle/VeloxShuffleReader.cc
b/cpp/velox/shuffle/VeloxShuffleReader.cc
index 0407be736a..3aba7cf0fc 100644
--- a/cpp/velox/shuffle/VeloxShuffleReader.cc
+++ b/cpp/velox/shuffle/VeloxShuffleReader.cc
@@ -15,13 +15,13 @@
* limitations under the License.
*/
-#include "VeloxShuffleReader.h"
-#include "GlutenByteStream.h"
+#include "shuffle/VeloxShuffleReader.h"
#include <arrow/array/array_binary.h>
#include <arrow/io/buffered.h>
#include "memory/VeloxColumnarBatch.h"
+#include "shuffle/GlutenByteStream.h"
#include "shuffle/Payload.h"
#include "shuffle/Utils.h"
#include "utils/Common.h"
@@ -576,7 +576,7 @@ std::shared_ptr<ColumnarBatch>
VeloxRssSortShuffleReaderDeserializer::next() {
return std::make_shared<VeloxColumnarBatch>(std::move(rowVector));
}
-VeloxColumnarBatchDeserializerFactory::VeloxColumnarBatchDeserializerFactory(
+VeloxShuffleReaderDeserializerFactory::VeloxShuffleReaderDeserializerFactory(
const std::shared_ptr<arrow::Schema>& schema,
const std::shared_ptr<arrow::util::Codec>& codec,
const facebook::velox::common::CompressionKind veloxCompressionType,
@@ -598,7 +598,7 @@
VeloxColumnarBatchDeserializerFactory::VeloxColumnarBatchDeserializerFactory(
initFromSchema();
}
-std::unique_ptr<ColumnarBatchIterator>
VeloxColumnarBatchDeserializerFactory::createDeserializer(
+std::unique_ptr<ColumnarBatchIterator>
VeloxShuffleReaderDeserializerFactory::createDeserializer(
std::shared_ptr<arrow::io::InputStream> in) {
switch (shuffleWriterType_) {
case ShuffleWriterType::kHashShuffle:
@@ -635,23 +635,19 @@ std::unique_ptr<ColumnarBatchIterator>
VeloxColumnarBatchDeserializerFactory::cr
}
}
-arrow::MemoryPool* VeloxColumnarBatchDeserializerFactory::getPool() {
+arrow::MemoryPool* VeloxShuffleReaderDeserializerFactory::getPool() {
return memoryPool_;
}
-ShuffleWriterType
VeloxColumnarBatchDeserializerFactory::getShuffleWriterType() {
- return shuffleWriterType_;
-}
-
-int64_t VeloxColumnarBatchDeserializerFactory::getDecompressTime() {
+int64_t VeloxShuffleReaderDeserializerFactory::getDecompressTime() {
return decompressTime_;
}
-int64_t VeloxColumnarBatchDeserializerFactory::getDeserializeTime() {
+int64_t VeloxShuffleReaderDeserializerFactory::getDeserializeTime() {
return deserializeTime_;
}
-void VeloxColumnarBatchDeserializerFactory::initFromSchema() {
+void VeloxShuffleReaderDeserializerFactory::initFromSchema() {
GLUTEN_ASSIGN_OR_THROW(auto arrowColumnTypes,
toShuffleTypeId(schema_->fields()));
isValidityBuffer_.reserve(arrowColumnTypes.size());
for (size_t i = 0; i < arrowColumnTypes.size(); ++i) {
@@ -681,7 +677,23 @@ void
VeloxColumnarBatchDeserializerFactory::initFromSchema() {
}
}
-VeloxShuffleReader::VeloxShuffleReader(std::unique_ptr<DeserializerFactory>
factory)
- : ShuffleReader(std::move(factory)) {}
+VeloxShuffleReader::VeloxShuffleReader(std::unique_ptr<VeloxShuffleReaderDeserializerFactory>
factory)
+ : factory_(std::move(factory)) {}
+
+std::shared_ptr<ResultIterator>
VeloxShuffleReader::readStream(std::shared_ptr<arrow::io::InputStream> in) {
+ return std::make_shared<ResultIterator>(factory_->createDeserializer(in));
+}
+
+arrow::MemoryPool* VeloxShuffleReader::getPool() const {
+ return factory_->getPool();
+}
+
+int64_t VeloxShuffleReader::getDecompressTime() const {
+ return factory_->getDecompressTime();
+}
+
+int64_t VeloxShuffleReader::getDeserializeTime() const {
+ return factory_->getDeserializeTime();
+}
} // namespace gluten
diff --git a/cpp/velox/shuffle/VeloxShuffleReader.h
b/cpp/velox/shuffle/VeloxShuffleReader.h
index af35f97712..8ebdbf2bac 100644
--- a/cpp/velox/shuffle/VeloxShuffleReader.h
+++ b/cpp/velox/shuffle/VeloxShuffleReader.h
@@ -17,16 +17,14 @@
#pragma once
-#include "operators/serializer/VeloxColumnarBatchSerializer.h"
#include "shuffle/Payload.h"
#include "shuffle/ShuffleReader.h"
#include "shuffle/VeloxSortShuffleWriter.h"
-#include "utils/Timer.h"
+
+#include "velox/serializers/PrestoSerializer.h"
#include "velox/type/Type.h"
#include "velox/vector/ComplexVector.h"
-#include <velox/serializers/PrestoSerializer.h>
-
namespace gluten {
class VeloxHashShuffleReaderDeserializer final : public ColumnarBatchIterator {
@@ -134,9 +132,9 @@ class VeloxRssSortShuffleReaderDeserializer : public
ColumnarBatchIterator {
std::shared_ptr<VeloxInputStream> in_;
};
-class VeloxColumnarBatchDeserializerFactory : public DeserializerFactory {
+class VeloxShuffleReaderDeserializerFactory {
public:
- VeloxColumnarBatchDeserializerFactory(
+ VeloxShuffleReaderDeserializerFactory(
const std::shared_ptr<arrow::Schema>& schema,
const std::shared_ptr<arrow::util::Codec>& codec,
const facebook::velox::common::CompressionKind veloxCompressionType,
@@ -147,15 +145,13 @@ class VeloxColumnarBatchDeserializerFactory : public
DeserializerFactory {
std::shared_ptr<facebook::velox::memory::MemoryPool> veloxPool,
ShuffleWriterType shuffleWriterType);
- std::unique_ptr<ColumnarBatchIterator>
createDeserializer(std::shared_ptr<arrow::io::InputStream> in) override;
-
- arrow::MemoryPool* getPool() override;
+ std::unique_ptr<ColumnarBatchIterator>
createDeserializer(std::shared_ptr<arrow::io::InputStream> in);
- int64_t getDecompressTime() override;
+ arrow::MemoryPool* getPool();
- int64_t getDeserializeTime() override;
+ int64_t getDecompressTime();
- ShuffleWriterType getShuffleWriterType() override;
+ int64_t getDeserializeTime();
private:
void initFromSchema();
@@ -180,6 +176,17 @@ class VeloxColumnarBatchDeserializerFactory : public
DeserializerFactory {
class VeloxShuffleReader final : public ShuffleReader {
public:
- VeloxShuffleReader(std::unique_ptr<DeserializerFactory> factory);
+ VeloxShuffleReader(std::unique_ptr<VeloxShuffleReaderDeserializerFactory>
factory);
+
+ std::shared_ptr<ResultIterator>
readStream(std::shared_ptr<arrow::io::InputStream> in) override;
+
+ int64_t getDecompressTime() const override;
+
+ int64_t getDeserializeTime() const override;
+
+ arrow::MemoryPool* getPool() const override;
+
+ private:
+ std::unique_ptr<VeloxShuffleReaderDeserializerFactory> factory_;
};
} // namespace gluten
diff --git a/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
b/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
index 9331a72078..4fcab8f242 100644
--- a/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
+++ b/cpp/velox/utils/tests/VeloxShuffleWriterTestBase.h
@@ -360,7 +360,7 @@ class VeloxShuffleWriterTest : public
::testing::TestWithParam<ShuffleTestParams
facebook::velox::serializer::presto::PrestoVectorSerde::registerVectorSerde();
}
// Set batchSize to a large value to make all batches are merged by reader.
- auto deserializerFactory =
std::make_unique<gluten::VeloxColumnarBatchDeserializerFactory>(
+ auto deserializerFactory =
std::make_unique<gluten::VeloxShuffleReaderDeserializerFactory>(
schema,
std::move(codec),
veloxCompressionType,
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
