Unnest changes for moving the filter on right side of correlate to inside the unnest datasource

processing/src/main/java/org/apache/druid/query/UnnestDataSource.java

@@ -23,12 +23,16 @@
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.google.common.collect.ImmutableList;
 import org.apache.druid.java.util.common.IAE;
+import org.apache.druid.query.filter.DimFilter;
 import org.apache.druid.query.planning.DataSourceAnalysis;
 import org.apache.druid.segment.SegmentReference;
 import org.apache.druid.segment.UnnestSegmentReference;
 import org.apache.druid.segment.VirtualColumn;
 import org.apache.druid.utils.JvmUtils;
 
+
+import javax.annotation.Nullable;
+
 import java.util.List;
 import java.util.Objects;
 import java.util.Set;
@@ -48,22 +52,29 @@ public class UnnestDataSource implements DataSource
   private final DataSource base;
   private final VirtualColumn virtualColumn;
 
+  @Nullable
+  private final DimFilter unnestFilter;
+
   private UnnestDataSource(
       DataSource dataSource,
-      VirtualColumn virtualColumn
+      VirtualColumn virtualColumn,
+      DimFilter unnestFilter
   )
   {
     this.base = dataSource;
     this.virtualColumn = virtualColumn;
+    this.unnestFilter = unnestFilter;
   }
 
   @JsonCreator
   public static UnnestDataSource create(
       @JsonProperty("base") DataSource base,
-      @JsonProperty("virtualColumn") VirtualColumn virtualColumn
+      @JsonProperty("virtualColumn") VirtualColumn virtualColumn,
+      @Nullable @JsonProperty("unnestFilter") DimFilter unnestFilter
+
   )
   {
-    return new UnnestDataSource(base, virtualColumn);
+    return new UnnestDataSource(base, virtualColumn, unnestFilter);
   }
 
   @JsonProperty("base")
@@ -78,6 +89,12 @@ public VirtualColumn getVirtualColumn()
     return virtualColumn;
   }
 
+  @JsonProperty("unnestFilter")
+  public DimFilter getUnnestFilter()
+  {
+    return unnestFilter;
+  }
+
   @Override
   public Set<String> getTableNames()
   {
@@ -96,7 +113,8 @@ public DataSource withChildren(List<DataSource> children)
     if (children.size() != 1) {
       throw new IAE("Expected [1] child, got [%d]", children.size());
     }
-    return new UnnestDataSource(children.get(0), virtualColumn);
+
+    return new UnnestDataSource(children.get(0), virtualColumn, unnestFilter);
   }
 
   @Override
@@ -133,15 +151,16 @@ public Function<SegmentReference, SegmentReference> createSegmentMapFunction(
             baseSegment ->
                 new UnnestSegmentReference(
                     segmentMapFn.apply(baseSegment),
-                    virtualColumn
+                    virtualColumn,
+                    unnestFilter
                 )
     );
   }
 
   @Override
   public DataSource withUpdatedDataSource(DataSource newSource)
   {
-    return new UnnestDataSource(newSource, virtualColumn);
+    return new UnnestDataSource(newSource, virtualColumn, unnestFilter);
   }
 
   @Override
@@ -162,6 +181,17 @@ public DataSourceAnalysis getAnalysis()
     return current.getAnalysis();
   }
 
+
+  @Override
+  public String toString()
+  {
+    return "UnnestDataSource{" +
+           "base=" + base +
+           ", column='" + virtualColumn + '\'' +
+           ", unnestFilter='" + unnestFilter + '\'' +
+           '}';
+  }
+
   @Override
   public boolean equals(Object o)
   {
@@ -172,25 +202,17 @@ public boolean equals(Object o)
       return false;
     }
     UnnestDataSource that = (UnnestDataSource) o;
-    return virtualColumn.equals(that.virtualColumn)
-           && base.equals(that.base);
+    return base.equals(that.base) && virtualColumn.equals(that.virtualColumn) && Objects.equals(
+        unnestFilter,
+        that.unnestFilter
+    );
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(base, virtualColumn);
-  }
-
-  @Override
-  public String toString()
-  {
-    return "UnnestDataSource{" +
-           "base=" + base +
-           ", column='" + virtualColumn + '\'' +
-           '}';
+    return Objects.hash(base, virtualColumn, unnestFilter);
   }
-
 }
 
 

processing/src/main/java/org/apache/druid/segment/UnnestDimensionCursor.java

@@ -142,6 +142,9 @@ public void inspectRuntimeShape(RuntimeShapeInspector inspector)
               @Override
               public boolean matches()
               {
+                if (indexedIntsForCurrentRow == null) {
+                  return false;
+                }
                 if (indexedIntsForCurrentRow.size() <= 0) {
                   return false;
                 }

processing/src/main/java/org/apache/druid/segment/UnnestSegmentReference.java

@@ -21,6 +21,7 @@
 
 import org.apache.druid.java.util.common.io.Closer;
 import org.apache.druid.java.util.common.logger.Logger;
+import org.apache.druid.query.filter.DimFilter;
 import org.apache.druid.timeline.SegmentId;
 import org.apache.druid.utils.CloseableUtils;
 import org.joda.time.Interval;
@@ -41,14 +42,20 @@ public class UnnestSegmentReference implements SegmentReference
   private final SegmentReference baseSegment;
   private final VirtualColumn unnestColumn;
 
+  @Nullable
+  private final DimFilter unnestFilter;
+
+
 
   public UnnestSegmentReference(
       SegmentReference baseSegment,
-      VirtualColumn unnestColumn
+      VirtualColumn unnestColumn,
+      DimFilter unnestFilter
   )
   {
     this.baseSegment = baseSegment;
     this.unnestColumn = unnestColumn;
+    this.unnestFilter = unnestFilter;
   }
 
   @Override
@@ -100,7 +107,8 @@ public StorageAdapter asStorageAdapter()
   {
     return new UnnestStorageAdapter(
         baseSegment.asStorageAdapter(),
-        unnestColumn
+        unnestColumn,
+        unnestFilter
     );
   }