[native] Add query trace session properties to native session

prestodb · Oct 14, 2024 · 27eb666 · 27eb666
1 parent 3236937
commit 27eb666
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 10 deletions.
diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.cpp b/presto-native-execution/presto_cpp/main/SessionProperties.cpp
@@ -177,20 +177,20 @@ SessionProperties::SessionProperties() {
       std::to_string(c.rowNumberSpillEnabled()));
 
   addSessionProperty(
-      kNativeSpillerNumPartitionBits,
+      kSpillerNumPartitionBits,
       "none",
       TINYINT(),
       false,
       QueryConfig::kSpillNumPartitionBits,
-      std::to_string(c.spillNumPartitionBits())),
+      std::to_string(c.spillNumPartitionBits()));
 
-      addSessionProperty(
-          kTopNRowNumberSpillEnabled,
-          "Native Execution only. Enable topN row number spilling on native engine",
-          BOOLEAN(),
-          false,
-          QueryConfig::kTopNRowNumberSpillEnabled,
-          boolToString(c.topNRowNumberSpillEnabled()));
+  addSessionProperty(
+      kTopNRowNumberSpillEnabled,
+      "Native Execution only. Enable topN row number spilling on native engine",
+      BOOLEAN(),
+      false,
+      QueryConfig::kTopNRowNumberSpillEnabled,
+      boolToString(c.topNRowNumberSpillEnabled()));
 
   addSessionProperty(
       kValidateOutputFromOperators,
@@ -254,6 +254,48 @@ SessionProperties::SessionProperties() {
       QueryConfig::kSelectiveNimbleReaderEnabled,
       boolToString(c.selectiveNimbleReaderEnabled()));
 
+  addSessionProperty(
+      kQueryTraceEnabled,
+      "Enables query tracing.",
+      BOOLEAN(),
+      false,
+      QueryConfig::kQueryTraceEnabled,
+      boolToString(c.queryTraceEnabled()));
+
+  addSessionProperty(
+      kQueryTraceDir,
+      "Base dir of a query to store tracing data.",
+      VARCHAR(),
+      false,
+      QueryConfig::kQueryTraceDir,
+      c.queryTraceDir());
+
+  addSessionProperty(
+      kQueryTraceNodeIds,
+      "A comma-separated list of plan node ids whose input data will be traced."
+      " Empty string if only want to trace the query metadata.",
+      VARCHAR(),
+      false,
+      QueryConfig::kQueryTraceNodeIds,
+      c.queryTraceNodeIds());
+
+  addSessionProperty(
+      kQueryTraceMaxBytes,
+      "The max trace bytes limit. Tracing is disabled if zero.",
+      BIGINT(),
+      false,
+      QueryConfig::kQueryTraceMaxBytes,
+      std::to_string(c.queryTraceMaxBytes()));
+
+  addSessionProperty(
+      kQueryTraceTaskRegExp,
+      "The regexp of traced task id. We only enable trace on a task if its id"
+      " matches.",
+      BIGINT(),
+      false,
+      QueryConfig::kQueryTraceTaskRegExp,
+      c.queryTraceTaskRegExp());
+
   // If `legacy_timestamp` is true, the coordinator expects timestamp
   // conversions without a timezone to be converted to the user's
   // session_timezone.

diff --git a/presto-native-execution/presto_cpp/main/SessionProperties.h b/presto-native-execution/presto_cpp/main/SessionProperties.h
@@ -131,7 +131,7 @@ class SessionProperties {
   static constexpr const char* kJoinSpillPartitionBits =
       "native_join_spiller_partition_bits";
 
-  static constexpr const char* kNativeSpillerNumPartitionBits =
+  static constexpr const char* kSpillerNumPartitionBits =
       "native_spiller_num_partition_bits";
 
   /// Enable topN row number spilling on native engine.
@@ -179,6 +179,27 @@ class SessionProperties {
   static constexpr const char* kDriverCpuTimeSliceLimitMs =
       "driver_cpu_time_slice_limit_ms";
 
+  /// Enables query tracing.
+  static constexpr const char* kQueryTraceEnabled =
+      "native_query_trace_enabled";
+
+  /// Base dir of a query to store tracing data.
+  static constexpr const char* kQueryTraceDir = "native_query_trace_dir";
+
+  /// A comma-separated list of plan node ids whose input data will be traced.
+  /// Empty string if only want to trace the query metadata.
+  static constexpr const char* kQueryTraceNodeIds =
+      "native_query_trace_node_ids";
+
+  /// The max trace bytes limit. Tracing is disabled if zero.
+  static constexpr const char* kQueryTraceMaxBytes =
+      "native_query_trace_max_bytes";
+
+  /// The regexp of traced task id. We only enable trace on a task if its id
+  /// matches.
+  static constexpr const char* kQueryTraceTaskRegExp =
+      "native_query_trace_task_reg_exp";
+
   SessionProperties();
 
   const std::unordered_map<std::string, std::shared_ptr<SessionProperty>>&