f

kgyrtkirk · Jan 14, 2025 · 293a483 · 293a483
1 parent bf001d1
commit 293a483
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 11 deletions.
diff --git a/processing/src/main/java/org/apache/druid/query/DataSource.java b/processing/src/main/java/org/apache/druid/query/DataSource.java
@@ -101,14 +101,27 @@ public interface DataSource
    */
   boolean isConcrete();
 
+  static class SegmentMapConfig {
+
+    private Query query;
+    private AtomicLong cpuTimeAcc;
+
+    public SegmentMapConfig(Query query, AtomicLong cpuTimeAcc)
+    {
+      this.query = query;
+      this.cpuTimeAcc = cpuTimeAcc;
+    }
+
+    static SegmentMapConfig of(Query query, AtomicLong cpuTimeAcc)
+    {
+      return new SegmentMapConfig(query,cpuTimeAcc);
+    }
+  }
+
   /**
    * Returns a segment function on to how to segment should be modified.
-   *
-   * @param query      the input query
-   * @param cpuTimeAcc the cpu time accumulator
-   * @return the segment function
    */
-  Function<SegmentReference, SegmentReference> createSegmentMapFunction(Query query, AtomicLong cpuTimeAcc);
+  Function<SegmentReference, SegmentReference> createSegmentMapFunction(SegmentMapConfig req);
 
   /**
    * Returns an updated datasource based on the specified new source.

diff --git a/processing/src/main/java/org/apache/druid/query/FilteredDataSource.java b/processing/src/main/java/org/apache/druid/query/FilteredDataSource.java
@@ -144,14 +144,10 @@ public boolean isConcrete()
   }
 
   @Override
-  public Function<SegmentReference, SegmentReference> createSegmentMapFunction(
-      Query query,
-      AtomicLong cpuTimeAccumulator
-  )
+  public Function<SegmentReference, SegmentReference> createSegmentMapFunction(SegmentMapConfig req)
   {
     final Function<SegmentReference, SegmentReference> segmentMapFn = base.createSegmentMapFunction(
-        query,
-        cpuTimeAccumulator
+        req
     );
     return JvmUtils.safeAccumulateThreadCpuTime(
         cpuTimeAccumulator,