From ed2dd88964e333414d7224334eed8d743573b237 Mon Sep 17 00:00:00 2001 From: Chengcheng Jin Date: Tue, 31 Dec 2024 10:17:52 +0000 Subject: [PATCH] support query trace --- cpp/velox/compute/WholeStageResultIterator.cc | 12 +++++ cpp/velox/config/VeloxConfig.h | 19 ++++++++ .../org/apache/gluten/GlutenConfig.scala | 47 ++++++++++++++++++- 3 files changed, 77 insertions(+), 1 deletion(-) diff --git a/cpp/velox/compute/WholeStageResultIterator.cc b/cpp/velox/compute/WholeStageResultIterator.cc index c1ae6f2184583..3a53daac8bec1 100644 --- a/cpp/velox/compute/WholeStageResultIterator.cc +++ b/cpp/velox/compute/WholeStageResultIterator.cc @@ -558,6 +558,18 @@ std::unordered_map WholeStageResultIterator::getQueryC configs[velox::core::QueryConfig::kSparkLegacyDateFormatter] = "false"; } + const auto setIfExists = [&](const std::string& glutenKey, const std::string& veloxKey) { + const auto valueOptional = veloxCfg_->get(glutenKey); + if (valueOptional.hasValue()) { + configs[veloxKey] = valueOptional.value(); + } + }; + setIfExists(kQueryTraceEnabled, velox::core::QueryConfig::kQueryTraceEnabled); + setIfExists(kQueryTraceDir, velox::core::QueryConfig::kQueryTraceDir); + setIfExists(kQueryTraceNodeIds, velox::core::QueryConfig::kQueryTraceNodeIds); + setIfExists(kQueryTraceMaxBytes, velox::core::QueryConfig::kQueryTraceMaxBytes); + setIfExists(kQueryTraceTaskRegExp, velox::core::QueryConfig::kQueryTraceTaskRegExp); + setIfExists(kOpTraceDirectoryCreateConfig, velox::core::QueryConfig::kOpTraceDirectoryCreateConfig); } catch (const std::invalid_argument& err) { std::string errDetails = err.what(); throw std::runtime_error("Invalid conf arg: " + errDetails); diff --git a/cpp/velox/config/VeloxConfig.h b/cpp/velox/config/VeloxConfig.h index f882e72065b2c..84493cc469198 100644 --- a/cpp/velox/config/VeloxConfig.h +++ b/cpp/velox/config/VeloxConfig.h @@ -135,4 +135,23 @@ const uint32_t kGlogVerboseLevelDefault = 0; const uint32_t kGlogVerboseLevelMaximum = 99; const std::string kGlogSeverityLevel = "spark.gluten.sql.columnar.backend.velox.glogSeverityLevel"; const uint32_t kGlogSeverityLevelDefault = 1; + +// Query trace +/// Enable query tracing flag. +const std::string kQueryTraceEnabled = "spark.gluten.sql.columnar.backend.velox.queryTraceEnabled"; +/// Base dir of a query to store tracing data. +const std::string kQueryTraceDir = "spark.gluten.sql.columnar.backend.velox.queryTraceDir"; +/// A comma-separated list of plan node ids whose input data will be traced. +/// Empty string if only want to trace the query metadata. +const std::string kQueryTraceNodeIds = "spark.gluten.sql.columnar.backend.velox.queryTraceNodeIds"; +/// The max trace bytes limit. Tracing is disabled if zero. +const std::string kQueryTraceMaxBytes = "spark.gluten.sql.columnar.backend.velox.queryTraceMaxBytes"; +/// The regexp of traced task id. We only enable trace on a task if its id +/// matches. +const std::string kQueryTraceTaskRegExp = "spark.gluten.sql.columnar.backend.velox.queryTraceTaskRegExp"; +/// Config used to create operator trace directory. This config is provided to +/// underlying file system and the config is free form. The form should be +/// defined by the underlying file system. +const std::string kOpTraceDirectoryCreateConfig = + "spark.gluten.sql.columnar.backend.velox.opTraceDirectoryCreateConfig"; } // namespace gluten diff --git a/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala b/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala index e4f69e5748f1a..ab9b94eaec7c5 100644 --- a/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala +++ b/shims/common/src/main/scala/org/apache/gluten/GlutenConfig.scala @@ -735,7 +735,13 @@ object GlutenConfig { SPARK_GCS_STORAGE_ROOT_URL, SPARK_GCS_AUTH_TYPE, SPARK_GCS_AUTH_SERVICE_ACCOUNT_JSON_KEYFILE, - SPARK_REDACTION_REGEX + SPARK_REDACTION_REGEX, + QUERY_TRACE_ENABLED, + QUERY_TRACE_DIR, + QUERY_TRACE_NODE_IDS, + QUERY_TRACE_MAX_BYTES, + QUERY_TRACE_TASK_REG_EXP, + OP_TRACE_DIRECTORY_CREATE_CONFIG ) nativeConfMap.putAll(conf.filter(e => keys.contains(e._1)).asJava) @@ -2268,4 +2274,43 @@ object GlutenConfig { "Otherwise, broadcast build relation will use onheap memory.") .booleanConf .createWithDefault(false) + + val QUERY_TRACE_ENABLED = buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceEnabled") + .doc("Enable query tracing flag.") + .booleanConf + .createWithDefault(false) + + val QUERY_TRACE_DIR = buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceDir") + .doc("Base dir of a query to store tracing data.") + .stringConf + .createWithDefault("") + + val QUERY_TRACE_NODE_IDS = buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceNodeIds") + .doc( + "A comma-separated list of plan node ids whose input data will be traced. " + + "Empty string if only want to trace the query metadata.") + .stringConf + .createWithDefault("") + + val QUERY_TRACE_MAX_BYTES = + buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceMaxBytes") + .doc("The max trace bytes limit. Tracing is disabled if zero.") + .longConf + .createWithDefault(0) + + val QUERY_TRACE_TASK_REG_EXP = + buildConf("spark.gluten.sql.columnar.backend.velox.queryTraceTaskRegExp") + .doc("The regexp of traced task id. We only enable trace on a task if its id matches.") + .stringConf + .createWithDefault("") + + val OP_TRACE_DIRECTORY_CREATE_CONFIG = + buildConf("spark.gluten.sql.columnar.backend.velox.opTraceDirectoryCreateConfig") + .doc( + "Config used to create operator trace directory. This config is provided to" + + " underlying file system and the config is free form. The form should be " + + "defined by the underlying file system.") + .stringConf + .createWithDefault("") + }