somu-imply commented on code in PR #13085:
URL: https://github.com/apache/druid/pull/13085#discussion_r997159469


##########
extensions-core/multi-stage-query/src/main/java/org/apache/druid/msq/querykit/BaseLeafFrameProcessor.java:
##########
@@ -146,19 +147,31 @@ protected SegmentReference mapSegment(final Segment 
segment)
 
   private boolean initializeSegmentMapFn(final IntSet readableInputs)
   {
+    final AtomicLong cpuAccumulator = new AtomicLong();
     if (segmentMapFn != null) {
       return true;
     } else if (broadcastJoinHelper == null) {
       segmentMapFn = Function.identity();
       return true;
     } else {
-      final boolean retVal = 
broadcastJoinHelper.buildBroadcastTablesIncrementally(readableInputs);
-
-      if (retVal) {
-        segmentMapFn = broadcastJoinHelper.makeSegmentMapFn(query);
+      if (query.getDataSource() instanceof InputNumberDataSource) {
+        final boolean retVal = 
broadcastJoinHelper.buildBroadcastTablesIncrementally(readableInputs);
+        if (retVal) {
+          InputNumberDataSource inputNumberDataSource = 
(InputNumberDataSource) query.getDataSource();
+          // The InputNumberDataSource requires a BroadcastJoinHelper to be 
able to create its
+          // segment map function.  It would be a lot better if the 
InputNumberDataSource actually
+          // had a way to get that injected into it on its own, but the 
relationship between these objects
+          // was figured out during a refactor and using a setter here seemed 
like the least-bad way to
+          // make progress on the refactor without breaking functionality.  
Hopefully, some future
+          // developer will move this away from a setter.
+          inputNumberDataSource.setBroadcastJoinHelper(broadcastJoinHelper);
+          segmentMapFn = inputNumberDataSource.createSegmentMapFunction(query, 
cpuAccumulator);
+        }
+        return retVal;
+      } else {
+        segmentMapFn = Function.identity();

Review Comment:
   To continue with Abhishek's point if the parent datasource is a 
`JoinDataSource` with child data sources being `InputNumberDataSource`, this 
would basically set the map function to an identity function which is not 
right. One way is to first construct the data source using the 
broadcastJoinHelper and call createSegmentMapFunction of that. This would also 
remove the instanceof check we are doing. I'll discuss this with @imply-cheddar 
today. I believe the else part can be simplified as
   ```
   else {
         final boolean retVal = 
broadcastJoinHelper.buildBroadcastTablesIncrementally(readableInputs);
         if (retVal) {
           final DataSource dataSourceWithInlinedChannelData = 
broadcastJoinHelper.inlineChannelData(query.getDataSource());
           segmentMapFn = 
dataSourceWithInlinedChannelData.createSegmentMapFunction(query, 
cpuAccumulator);
         }
         else {
           segmentMapFn = Function.identity();
         }
         return true;
   }
   ```



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to