Add two new parameters for trace control (#24209)

Add two new parameters for trace control

presto-docs/src/main/sphinx/presto_cpp/properties-session.rst

@@ -307,11 +307,73 @@ Minimum number of rows to use prefix-sort.
 The default value has been derived using micro-benchmarking.
 
 ``native_op_trace_directory_create_config``
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 * **Type:** ``varchar``
 * **Default value:** ``""``
 
 Native Execution only. 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.
+underlying file system.
+
+``native_query_trace_enabled``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``boolean``
+* **Default value:** ``false``
+
+Enable query tracing. After enabled, trace data will be generated with query execution, and
+can be used by TraceReplayer. It needs to be used together with native_query_trace_node_ids,
+native_query_trace_max_bytes, native_query_trace_fragment_id, and native_query_trace_shard_id
+to match the task to be traced.
+
+
+``native_query_trace_dir``
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``""``
+
+The location to store the trace files.
+
+``native_query_trace_node_ids``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``""``
+
+A comma-separated list of plan node ids whose input data will be traced.
+Empty string if only want to trace the query metadata.
+
+``native_query_trace_max_bytes``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``integer``
+* **Default value:** ``0``
+
+The max trace bytes limit. Tracing is disabled if zero.
+
+``native_query_trace_fragment_id``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``.*``
+
+The fragment id to be traced. If not specified, all fragments will be matched.
+
+``native_query_trace_shard_id``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``.*``
+
+The shard id to be traced. If not specified, all shards will be matched.
+
+``native_query_trace_task_reg_exp``
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+* **Type:** ``varchar``
+* **Default value:** ``""``
+
+The regular expression to match a task for tracing. It will be deprecated if there is
+no issue with native_query_trace_fragment_id and native_query_trace_shard_id.

presto-main/src/main/java/com/facebook/presto/sessionpropertyproviders/NativeWorkerSessionPropertyProvider.java

@@ -61,6 +61,8 @@ public class NativeWorkerSessionPropertyProvider
     public static final String NATIVE_QUERY_TRACE_NODE_IDS = "native_query_trace_node_ids";
     public static final String NATIVE_QUERY_TRACE_MAX_BYTES = "native_query_trace_max_bytes";
     public static final String NATIVE_QUERY_TRACE_REG_EXP = "native_query_trace_task_reg_exp";
+    public static final String NATIVE_QUERY_TRACE_FRAGMENT_ID = "native_query_trace_fragment_id";
+    public static final String NATIVE_QUERY_TRACE_SHARD_ID = "native_query_trace_shard_id";
     public static final String NATIVE_MAX_LOCAL_EXCHANGE_PARTITION_COUNT = "native_max_local_exchange_partition_count";
     public static final String NATIVE_SPILL_PREFIXSORT_ENABLED = "native_spill_prefixsort_enabled";
     public static final String NATIVE_PREFIXSORT_NORMALIZED_KEY_MAX_BYTES = "native_prefixsort_normalized_key_max_bytes";
@@ -231,6 +233,14 @@ public NativeWorkerSessionPropertyProvider(FeaturesConfig featuresConfig)
                         "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.",
                         "",
                         !nativeExecution),
+                stringProperty(NATIVE_QUERY_TRACE_FRAGMENT_ID,
+                            "The fragment id of the traced task.",
+                        "",
+                        !nativeExecution),
+                stringProperty(NATIVE_QUERY_TRACE_SHARD_ID,
+                        "The shard id of the traced task.",
+                        "",
+                        !nativeExecution),
                 longProperty(NATIVE_MAX_OUTPUT_BUFFER_SIZE,
                         "The maximum size in bytes for the task's buffered output. The buffer is shared among all drivers.",
                         200L << 20,

presto-native-execution/presto_cpp/main/QueryContextManager.cpp

@@ -217,9 +217,17 @@ QueryContextManager::toVeloxConfigs(
   // Use base velox query config as the starting point and add Presto session
   // properties on top of it.
   auto configs = BaseVeloxQueryConfig::instance()->values();
+  std::optional<std::string> traceFragmentId;
+  std::optional<std::string> traceShardId;
   for (const auto& it : session.systemProperties) {
-    configs[sessionProperties_.toVeloxConfig(it.first)] = it.second;
-    sessionProperties_.updateVeloxConfig(it.first, it.second);
+    if (it.first == SessionProperties::kQueryTraceFragmentId) {
+      traceFragmentId = it.second;
+    } else if (it.first == SessionProperties::kQueryTraceShardId) {
+      traceShardId = it.second;
+    } else {
+      configs[sessionProperties_.toVeloxConfig(it.first)] = it.second;
+      sessionProperties_.updateVeloxConfig(it.first, it.second);
+    }
   }
 
   // If there's a timeZoneKey, convert to timezone name and add to the
@@ -229,6 +237,16 @@ QueryContextManager::toVeloxConfigs(
         velox::core::QueryConfig::kSessionTimezone,
         velox::tz::getTimeZoneName(session.timeZoneKey));
   }
+
+  // Construct query tracing regex and pass to Velox config.
+  // It replaces the given native_query_trace_task_reg_exp if also set.
+  if (traceFragmentId.has_value() || traceShardId.has_value()) {
+    configs.emplace(
+        velox::core::QueryConfig::kQueryTraceTaskRegExp,
+        ".*\\." + traceFragmentId.value_or(".*") + "\\..*\\." +
+            traceShardId.value_or(".*") + "\\..*");
+  }
+
   updateFromSystemConfigs(configs);
   return configs;
 }

presto-native-execution/presto_cpp/main/SessionProperties.h

@@ -213,6 +213,16 @@ class SessionProperties {
   static constexpr const char* kOpTraceDirectoryCreateConfig =
       "native_op_trace_directory_create_config";
 
+  /// The fragment id of the traced task. Used to construct
+  /// the regular expression for matching
+  static constexpr const char* kQueryTraceFragmentId =
+      "native_query_trace_fragment_id";
+
+  /// The shard id of the traced task. Used to construct
+  /// the regular expression for matching
+  static constexpr const char* kQueryTraceShardId =
+      "native_query_trace_shard_id";
+
   /// The maximum size in bytes for the task's buffered output. The buffer is
   /// shared among all drivers.
   static constexpr const char* kMaxOutputBufferSize =