Additional native query tests for unnest datasource (#13554)

Native tests for the unnest datasource.
apache · Jan 25, 2023 · 17c0167 · 17c0167
1 parent 00cee32
commit 17c0167
Show file tree

Hide file tree

Showing 5 changed files with 1,821 additions and 6 deletions.
diff --git a/processing/src/test/java/org/apache/druid/query/QueryRunnerTestHelper.java b/processing/src/test/java/org/apache/druid/query/QueryRunnerTestHelper.java
@@ -63,6 +63,7 @@
 import org.apache.druid.segment.QueryableIndexSegment;
 import org.apache.druid.segment.ReferenceCountingSegment;
 import org.apache.druid.segment.Segment;
+import org.apache.druid.segment.SegmentReference;
 import org.apache.druid.segment.TestIndex;
 import org.apache.druid.segment.incremental.IncrementalIndex;
 import org.apache.druid.timeline.SegmentId;
@@ -79,6 +80,7 @@
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
@@ -100,6 +102,13 @@ public class QueryRunnerTestHelper
             .collect(Collectors.toList())
   );
 
+  public static final DataSource UNNEST_DATA_SOURCE = UnnestDataSource.create(
+      new TableDataSource(QueryRunnerTestHelper.DATA_SOURCE),
+      QueryRunnerTestHelper.PLACEMENTISH_DIMENSION,
+      QueryRunnerTestHelper.PLACEMENTISH_DIMENSION_UNNEST,
+      null
+  );
+
   public static final Granularity DAY_GRAN = Granularities.DAY;
   public static final Granularity ALL_GRAN = Granularities.ALL;
   public static final Granularity MONTH_GRAN = Granularities.MONTH;
@@ -109,6 +118,7 @@ public class QueryRunnerTestHelper
   public static final String PLACEMENT_DIMENSION = "placement";
   public static final String PLACEMENTISH_DIMENSION = "placementish";
   public static final String PARTIAL_NULL_DIMENSION = "partial_null_column";
+  public static final String PLACEMENTISH_DIMENSION_UNNEST = "placementish_unnest";
 
   public static final List<String> DIMENSIONS = Lists.newArrayList(
       MARKET_DIMENSION,
@@ -353,6 +363,7 @@ public static boolean isTestRunnerVectorizable(QueryRunner runner)
     return !("rtIndex".equals(runnerName) || "noRollupRtIndex".equals(runnerName));
   }
 
+
   public static <T, QueryType extends Query<T>> List<QueryRunner<T>> makeQueryRunners(
       QueryRunnerFactory<T, QueryType> factory
   )
@@ -395,6 +406,7 @@ public static Collection<?> makeUnionQueryRunners(QueryRunnerFactory factory)
     );
   }
 
+
   public static <T, QueryType extends Query<T>> QueryRunner<T> makeQueryRunner(
       QueryRunnerFactory<T, QueryType> factory,
       String resourceFileName,
@@ -467,6 +479,22 @@ public String toString()
     };
   }
 
+
+  public static <T, QueryType extends Query<T>> QueryRunner<T> makeQueryRunnerWithSegmentMapFn(
+      QueryRunnerFactory<T, QueryType> factory,
+      Segment adapter,
+      Query<T> query,
+      final String runnerName
+  )
+  {
+    final DataSource base = query.getDataSource();
+
+    final SegmentReference segmentReference = base.createSegmentMapFunction(query, new AtomicLong())
+                                                        .apply(ReferenceCountingSegment.wrapRootGenerationSegment(
+                                                            adapter));
+    return makeQueryRunner(factory, segmentReference, runnerName);
+  }
+
   public static <T> QueryRunner<T> makeFilteringQueryRunner(
       final VersionedIntervalTimeline<String, ReferenceCountingSegment> timeline,
       final QueryRunnerFactory<T, Query<T>> factory