Skip to content

Commit c67b735

Browse files
infvginfvg
committed
[VL] Refactor Gluten to use upstream Velox
1 parent 473a5d2 commit c67b735

17 files changed

Lines changed: 146 additions & 63 deletions

File tree

‎backends-velox/src-iceberg/test/scala/org/apache/gluten/execution/enhanced/VeloxIcebergSuite.scala‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -327,8 +327,7 @@ class VeloxIcebergSuite extends IcebergSuite {
327327
val lastExecId = statusStore.executionsList().last.executionId
328328
val executionMetrics = statusStore.executionMetrics(lastExecId)
329329

330-
// TODO: fix https://github.com/apache/gluten/issues/11510
331-
assert(executionMetrics(metrics("numWrittenFiles").id).toLong == 0)
330+
assert(executionMetrics(metrics("numWrittenFiles").id).toLong == 1)
332331
}
333332
}
334333

‎cpp/velox/CMakeLists.txt‎

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -230,10 +230,8 @@ if(ENABLE_GPU)
230230
memory/GpuBufferColumnarBatch.cc)
231231
endif()
232232

233-
if(ENABLE_ENHANCED_FEATURES)
234-
list(APPEND VELOX_SRCS compute/iceberg/IcebergFormat.cc
235-
compute/iceberg/IcebergWriter.cc)
236-
endif()
233+
list(APPEND VELOX_SRCS compute/iceberg/IcebergFormat.cc
234+
compute/iceberg/IcebergWriter.cc)
237235

238236
if(BUILD_TESTS OR BUILD_BENCHMARKS)
239237
set(BUILD_TEST_UTILS ON)

‎cpp/velox/compute/VeloxBackend.cc‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,7 @@
5555
#include "velox/connectors/hive/BufferedInputBuilder.h"
5656
#include "velox/connectors/hive/HiveConnector.h"
5757
#include "velox/connectors/hive/HiveDataSource.h"
58+
#include "velox/connectors/hive/iceberg/IcebergConnector.h"
5859
#include "velox/connectors/hive/storage_adapters/abfs/RegisterAbfsFileSystem.h" // @manual
5960
#include "velox/connectors/hive/storage_adapters/gcs/RegisterGcsFileSystem.h" // @manual
6061
#include "velox/connectors/hive/storage_adapters/hdfs/HdfsFileSystem.h"
@@ -378,6 +379,13 @@ std::shared_ptr<facebook::velox::connector::Connector> VeloxBackend::createDelta
378379
return std::make_shared<delta::DeltaConnector>(connectorId, hiveConnectorConfig_, ioExecutor);
379380
}
380381

382+
std::shared_ptr<facebook::velox::connector::Connector> VeloxBackend::createIcebergConnector(
383+
const std::string& connectorId,
384+
folly::Executor* ioExecutor) const {
385+
return std::make_shared<velox::connector::hive::iceberg::IcebergConnector>(
386+
connectorId, hiveConnectorConfig_, ioExecutor);
387+
}
388+
381389
std::shared_ptr<facebook::velox::connector::Connector> VeloxBackend::createValueStreamConnector(
382390
const std::string& connectorId,
383391
bool dynamicFilterEnabled) const {

‎cpp/velox/compute/VeloxBackend.h‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,10 @@ class VeloxBackend {
7474
const std::string& connectorId,
7575
folly::Executor* ioExecutor) const;
7676

77+
std::shared_ptr<facebook::velox::connector::Connector> createIcebergConnector(
78+
const std::string& connectorId,
79+
folly::Executor* ioExecutor) const;
80+
7781
std::shared_ptr<facebook::velox::connector::Connector> createDeltaConnector(
7882
const std::string& connectorId,
7983
folly::Executor* ioExecutor) const;

‎cpp/velox/compute/VeloxConnectorIds.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,10 +23,12 @@ namespace gluten {
2323

2424
struct VeloxConnectorIds {
2525
std::string hive;
26+
std::string iceberg;
2627
std::string delta;
2728
std::string iterator;
2829
std::string cudfHive;
2930
bool hiveRegistered{false};
31+
bool icebergRegistered{false};
3032
bool deltaRegistered{false};
3133
bool iteratorRegistered{false};
3234
bool cudfHiveRegistered{false};

‎cpp/velox/compute/VeloxRuntime.cc‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@ std::string makeScopedConnectorId(const std::string& base, uint64_t runtimeId) {
227227
VeloxConnectorIds makeScopedConnectorIds(uint64_t runtimeId) {
228228
return VeloxConnectorIds{
229229
.hive = makeScopedConnectorId(kHiveConnectorId, runtimeId),
230+
.iceberg = makeScopedConnectorId(kIcebergConnectorId, runtimeId),
230231
.delta = makeScopedConnectorId(delta::DeltaConnectorFactory::kDeltaConnectorName, runtimeId),
231232
.iterator = makeScopedConnectorId(kIteratorConnectorId, runtimeId),
232233
.cudfHive = makeScopedConnectorId(kCudfHiveConnectorId, runtimeId)};
@@ -290,6 +291,14 @@ void VeloxRuntime::registerConnectors() {
290291
velox::connector::hasConnector(connectorIds_.hive),
291292
"Scoped hive connector not found after registration: " + connectorIds_.hive);
292293

294+
connectorIds_.icebergRegistered =
295+
velox::connector::registerConnector(backend->createIcebergConnector(connectorIds_.iceberg, ioExecutor_.get()));
296+
GLUTEN_CHECK(
297+
connectorIds_.icebergRegistered, "Failed to register scoped Iceberg connector: " + connectorIds_.iceberg);
298+
GLUTEN_CHECK(
299+
velox::connector::hasConnector(connectorIds_.iceberg),
300+
"Scoped Iceberg connector not found after registration: " + connectorIds_.iceberg);
301+
293302
connectorIds_.deltaRegistered =
294303
velox::connector::registerConnector(backend->createDeltaConnector(connectorIds_.delta, ioExecutor_.get()));
295304
GLUTEN_CHECK(connectorIds_.deltaRegistered, "Failed to register scoped delta connector: " + connectorIds_.delta);
@@ -340,6 +349,10 @@ void VeloxRuntime::unregisterConnectors() {
340349
velox::connector::unregisterConnector(connectorIds_.hive);
341350
connectorIds_.hiveRegistered = false;
342351
}
352+
if (connectorIds_.icebergRegistered) {
353+
velox::connector::unregisterConnector(connectorIds_.iceberg);
354+
connectorIds_.icebergRegistered = false;
355+
}
343356
}
344357

345358
void VeloxRuntime::parsePlan(const uint8_t* data, int32_t size) {
@@ -518,7 +531,6 @@ std::shared_ptr<RowToColumnarConverter> VeloxRuntime::createRow2ColumnarConverte
518531
return std::make_shared<VeloxRowToColumnarConverter>(cSchema, veloxPool);
519532
}
520533

521-
#ifdef GLUTEN_ENABLE_ENHANCED_FEATURES
522534
std::shared_ptr<IcebergWriter> VeloxRuntime::createIcebergWriter(
523535
RowTypePtr rowType,
524536
int32_t format,
@@ -546,7 +558,6 @@ std::shared_ptr<IcebergWriter> VeloxRuntime::createIcebergWriter(
546558
veloxPool,
547559
connectorPool);
548560
}
549-
#endif
550561

551562
std::shared_ptr<ShuffleWriter> VeloxRuntime::createShuffleWriter(
552563
int32_t numPartitions,

‎cpp/velox/compute/VeloxRuntime.h‎

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -20,20 +20,15 @@
2020
#include "WholeStageResultIterator.h"
2121
#include "compute/Runtime.h"
2222
#include "compute/VeloxConnectorIds.h"
23-
#ifdef GLUTEN_ENABLE_ENHANCED_FEATURES
2423
#include "iceberg/IcebergWriter.h"
25-
#endif
2624
#include <folly/Executor.h>
2725
#include "memory/VeloxMemoryManager.h"
2826
#include "operators/serializer/VeloxColumnarBatchSerializer.h"
2927
#include "operators/serializer/VeloxColumnarToRowConverter.h"
3028
#include "operators/writer/VeloxParquetDataSource.h"
3129
#include "shuffle/ShuffleReader.h"
3230
#include "shuffle/ShuffleWriter.h"
33-
34-
#ifdef GLUTEN_ENABLE_ENHANCED_FEATURES
3531
#include "IcebergNestedField.pb.h"
36-
#endif
3732

3833
namespace gluten {
3934

@@ -77,7 +72,6 @@ class VeloxRuntime final : public Runtime {
7772

7873
std::shared_ptr<RowToColumnarConverter> createRow2ColumnarConverter(struct ArrowSchema* cSchema) override;
7974

80-
#ifdef GLUTEN_ENABLE_ENHANCED_FEATURES
8175
std::shared_ptr<IcebergWriter> createIcebergWriter(
8276
RowTypePtr rowType,
8377
int32_t format,
@@ -89,7 +83,6 @@ class VeloxRuntime final : public Runtime {
8983
std::shared_ptr<const facebook::velox::connector::hive::iceberg::IcebergPartitionSpec> spec,
9084
const gluten::IcebergNestedField& protoField,
9185
const std::unordered_map<std::string, std::string>& sparkConfs);
92-
#endif
9386

9487
std::shared_ptr<ShuffleWriter> createShuffleWriter(
9588
int numPartitions,

‎cpp/velox/compute/WholeStageResultIterator.cc‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -202,7 +202,7 @@ WholeStageResultIterator::WholeStageResultIterator(
202202
std::unordered_map<std::string, std::string> customSplitInfo{{"table_format", "hive-iceberg"}};
203203
auto deleteFiles = icebergSplitInfo->deleteFilesVec[idx];
204204
split = std::make_shared<velox::connector::hive::iceberg::HiveIcebergSplit>(
205-
connectorIds_.hive,
205+
connectorIds_.iceberg,
206206
paths[idx],
207207
format,
208208
starts[idx],
@@ -284,6 +284,7 @@ std::shared_ptr<velox::core::QueryCtx> WholeStageResultIterator::createNewVeloxQ
284284
std::unordered_map<std::string, std::shared_ptr<velox::config::ConfigBase>> connectorConfigs;
285285
auto hiveSessionConfig = createHiveConnectorSessionConfig(veloxCfg_);
286286
connectorConfigs[connectorIds_.hive] = hiveSessionConfig;
287+
connectorConfigs[connectorIds_.iceberg] = hiveSessionConfig;
287288
connectorConfigs[connectorIds_.delta] = hiveSessionConfig;
288289
connectorConfigs[connectorIds_.iterator] = hiveSessionConfig;
289290
#ifdef GLUTEN_ENABLE_GPU

‎cpp/velox/compute/iceberg/IcebergWriter.cc‎

Lines changed: 31 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,10 @@
1717

1818
#include "IcebergWriter.h"
1919

20+
#include "IcebergNestedField.pb.h"
2021
#include "IcebergPartitionSpec.pb.h"
2122
#include "compute/ProtobufUtils.h"
23+
#include "compute/VeloxBackend.h"
2224
#include "compute/iceberg/IcebergFormat.h"
2325
#include "config/VeloxConfig.h"
2426
#include "utils/ConfigExtractor.h"
@@ -99,9 +101,9 @@ class GlutenIcebergFileNameGenerator : public connector::hive::FileNameGenerator
99101
mutable int32_t fileCount_;
100102
};
101103

102-
iceberg::IcebergNestedField convertToIcebergNestedField(const gluten::IcebergNestedField& protoField) {
103-
IcebergNestedField result;
104-
result.id = protoField.id();
104+
parquet::ParquetFieldId convertToIcebergNestedField(const gluten::IcebergNestedField& protoField) {
105+
parquet::ParquetFieldId result;
106+
result.fieldId = protoField.id();
105107

106108
// Recursively convert children
107109
result.children.reserve(protoField.children_size());
@@ -121,7 +123,7 @@ std::shared_ptr<IcebergInsertTableHandle> createIcebergInsertTableHandle(
121123
int64_t taskId,
122124
const std::string& operationId,
123125
std::shared_ptr<const IcebergPartitionSpec> spec,
124-
const iceberg::IcebergNestedField& nestedField,
126+
const parquet::ParquetFieldId& nestedField,
125127
facebook::velox::memory::MemoryPool* pool) {
126128
std::vector<std::shared_ptr<const iceberg::IcebergColumnHandle>> columnHandles;
127129

@@ -139,14 +141,12 @@ std::shared_ptr<IcebergInsertTableHandle> createIcebergInsertTableHandle(
139141
columnNames.at(i),
140142
connector::hive::HiveColumnHandle::ColumnType::kPartitionKey,
141143
columnTypes.at(i),
142-
columnTypes.at(i),
143144
nestedField.children[i]));
144145
} else {
145146
columnHandles.push_back(std::make_shared<iceberg::IcebergColumnHandle>(
146147
columnNames.at(i),
147148
connector::hive::HiveColumnHandle::ColumnType::kRegular,
148149
columnTypes.at(i),
149-
columnTypes.at(i),
150150
nestedField.children[i]));
151151
}
152152
}
@@ -157,18 +157,10 @@ std::shared_ptr<IcebergInsertTableHandle> createIcebergInsertTableHandle(
157157
std::shared_ptr<const connector::hive::LocationHandle> locationHandle =
158158
std::make_shared<connector::hive::LocationHandle>(
159159
outputDirectoryPath, outputDirectoryPath, connector::hive::LocationHandle::TableType::kExisting);
160-
const std::vector<IcebergSortingColumn> sortedBy;
161160
const std::unordered_map<std::string, std::string> serdeParameters;
161+
auto writeKind = connector::hive::iceberg::IcebergInsertTableHandle::WriteKind::kData;
162162
return std::make_shared<connector::hive::iceberg::IcebergInsertTableHandle>(
163-
columnHandles,
164-
locationHandle,
165-
spec,
166-
pool,
167-
fileFormat,
168-
sortedBy,
169-
compressionKind,
170-
serdeParameters,
171-
fileNameGenerator);
163+
columnHandles, locationHandle, fileFormat, spec, compressionKind, serdeParameters, writeKind, fileNameGenerator);
172164
}
173165

174166
} // namespace
@@ -200,20 +192,36 @@ IcebergWriter::IcebergWriter(
200192
connectorSessionProperties_ = createHiveConnectorSessionConfig(veloxCfg);
201193
connectorConfig_ =
202194
std::make_shared<facebook::velox::connector::hive::HiveConfig>(createHiveConnectorConfig(veloxCfg));
195+
std::unordered_map<std::string, std::shared_ptr<facebook::velox::config::ConfigBase>> connectorConfigs;
196+
connectorConfigs[kHiveConnectorId] = connectorSessionProperties_;
197+
auto queryConfigBase =
198+
std::make_shared<facebook::velox::config::ConfigBase>(std::unordered_map<std::string, std::string>(sparkConfs));
199+
queryCtx_ = facebook::velox::core::QueryCtx::create(
200+
nullptr,
201+
facebook::velox::core::QueryConfig{facebook::velox::core::QueryConfig::ConfigTag{}, queryConfigBase},
202+
connectorConfigs,
203+
nullptr, // cache
204+
pool_,
205+
nullptr, // spillExecutor
206+
"IcebergWriter");
207+
208+
auto expressionEvaluator =
209+
std::make_unique<facebook::velox::exec::SimpleExpressionEvaluator>(queryCtx_.get(), pool_.get());
210+
203211
connectorQueryCtx_ = std::make_unique<connector::ConnectorQueryCtx>(
204212
pool_.get(),
205213
connectorPool_.get(),
206214
connectorSessionProperties_.get(),
207215
nullptr,
208216
common::PrefixSortConfig(),
209-
nullptr,
217+
std::move(expressionEvaluator),
210218
nullptr,
211219
"query.IcebergDataSink",
212220
"task.IcebergDataSink",
213221
"planNodeId.IcebergDataSink",
214222
0,
215223
"");
216-
224+
auto icebergConfig = std::make_shared<facebook::velox::connector::hive::iceberg::IcebergConfig>(veloxCfg);
217225
dataSink_ = std::make_unique<IcebergDataSink>(
218226
rowType_,
219227
createIcebergInsertTableHandle(
@@ -229,7 +237,8 @@ IcebergWriter::IcebergWriter(
229237
pool_.get()),
230238
connectorQueryCtx_.get(),
231239
facebook::velox::connector::CommitStrategy::kNoCommit,
232-
connectorConfig_);
240+
connectorConfig_,
241+
icebergConfig);
233242
}
234243

235244
void IcebergWriter::write(const VeloxColumnarBatch& batch) {
@@ -238,6 +247,9 @@ void IcebergWriter::write(const VeloxColumnarBatch& batch) {
238247

239248
if (inputRowType->size() != rowType_->size()) {
240249
const auto& children = inputRowVector->children();
250+
251+
VELOX_CHECK_GE(children.size(), 1 + rowType_->size());
252+
241253
std::vector<VectorPtr> dataColumns(children.begin() + 1, children.begin() + 1 + rowType_->size());
242254

243255
auto filteredRowVector = std::make_shared<RowVector>(

‎cpp/velox/compute/iceberg/IcebergWriter.h‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -47,7 +47,7 @@ class IcebergWriter {
4747
int64_t taskId,
4848
const std::string& operationId,
4949
std::shared_ptr<const facebook::velox::connector::hive::iceberg::IcebergPartitionSpec> spec,
50-
const gluten::IcebergNestedField& field,
50+
const IcebergNestedField& field,
5151
const std::unordered_map<std::string, std::string>& sparkConfs,
5252
std::shared_ptr<facebook::velox::memory::MemoryPool> memoryPool,
5353
std::shared_ptr<facebook::velox::memory::MemoryPool> connectorPool);
@@ -60,7 +60,7 @@ class IcebergWriter {
6060

6161
private:
6262
facebook::velox::RowTypePtr rowType_;
63-
const facebook::velox::connector::hive::iceberg::IcebergNestedField field_;
63+
const facebook::velox::parquet::ParquetFieldId field_;
6464
int32_t partitionId_;
6565
int64_t taskId_;
6666
std::string operationId_;
@@ -69,6 +69,7 @@ class IcebergWriter {
6969
std::shared_ptr<facebook::velox::connector::hive::HiveConfig> connectorConfig_;
7070
std::shared_ptr<facebook::velox::config::ConfigBase> connectorSessionProperties_;
7171

72+
std::shared_ptr<facebook::velox::core::QueryCtx> queryCtx_;
7273
std::unique_ptr<facebook::velox::connector::ConnectorQueryCtx> connectorQueryCtx_;
7374

7475
std::unique_ptr<facebook::velox::connector::hive::iceberg::IcebergDataSink> dataSink_;

0 commit comments

Comments
 (0)