Any virtual column on "__time" should be a pre-join virtual column (#10451)

* Virtual column on __time should be in pre-join

* Add unit test

processing/src/main/java/org/apache/druid/segment/join/HashJoinSegmentStorageAdapter.java

@@ -34,6 +34,7 @@ import org.apache.druid.segment.StorageAdapter;
 import org.apache.druid.segment.VirtualColumn;
 import org.apache.druid.segment.VirtualColumns;
 import org.apache.druid.segment.column.ColumnCapabilities;
+import org.apache.druid.segment.column.ColumnHolder;
 import org.apache.druid.segment.data.Indexed;
 import org.apache.druid.segment.data.ListIndexed;
 import org.apache.druid.segment.join.filter.JoinFilterAnalyzer;
@@ -305,6 +306,7 @@ public class HashJoinSegmentStorageAdapter implements StorageAdapter
   )
   {
     final Set<String> baseColumns = new HashSet<>();
+    baseColumns.add(ColumnHolder.TIME_COLUMN_NAME);
     Iterables.addAll(baseColumns, baseAdapter.getAvailableDimensions());
     Iterables.addAll(baseColumns, baseAdapter.getAvailableMetrics());
 

processing/src/test/java/org/apache/druid/segment/join/BaseHashJoinSegmentStorageAdapterTest.java

@@ -27,7 +27,9 @@ import org.apache.druid.query.QueryContexts;
 import org.apache.druid.query.filter.Filter;
 import org.apache.druid.query.lookup.LookupExtractor;
 import org.apache.druid.segment.QueryableIndexSegment;
+import org.apache.druid.segment.VirtualColumn;
 import org.apache.druid.segment.VirtualColumns;
+import org.apache.druid.segment.column.ValueType;
 import org.apache.druid.segment.join.filter.JoinFilterAnalyzer;
 import org.apache.druid.segment.join.filter.JoinFilterPreAnalysis;
 import org.apache.druid.segment.join.filter.JoinFilterPreAnalysisKey;
@@ -242,4 +244,19 @@ public class BaseHashJoinSegmentStorageAdapterTest
         )
     );
   }
+
+  protected VirtualColumn makeExpressionVirtualColumn(String expression)
+  {
+    return makeExpressionVirtualColumn(expression, "virtual");
+  }
+
+  protected VirtualColumn makeExpressionVirtualColumn(String expression, String columnName)
+  {
+    return new ExpressionVirtualColumn(
+        columnName,
+        expression,
+        ValueType.STRING,
+        ExprMacroTable.nil()
+    );
+  }
 }

processing/src/test/java/org/apache/druid/segment/join/HashJoinSegmentStorageAdapterTest.java

@@ -31,6 +31,7 @@ import org.apache.druid.query.filter.ExpressionDimFilter;
 import org.apache.druid.query.filter.Filter;
 import org.apache.druid.query.filter.OrDimFilter;
 import org.apache.druid.query.filter.SelectorDimFilter;
+import org.apache.druid.segment.VirtualColumn;
 import org.apache.druid.segment.VirtualColumns;
 import org.apache.druid.segment.column.ColumnCapabilities;
 import org.apache.druid.segment.column.ValueType;
@@ -38,10 +39,10 @@ import org.apache.druid.segment.filter.SelectorFilter;
 import org.apache.druid.segment.join.filter.JoinFilterPreAnalysis;
 import org.apache.druid.segment.join.lookup.LookupJoinable;
 import org.apache.druid.segment.join.table.IndexedTableJoinable;
-import org.apache.druid.segment.virtual.ExpressionVirtualColumn;
 import org.junit.Assert;
 import org.junit.Test;
 
+import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
 
@@ -1502,12 +1503,7 @@ public class HashJoinSegmentStorageAdapterTest extends BaseHashJoinSegmentStorag
 
     VirtualColumns virtualColumns = VirtualColumns.create(
         Collections.singletonList(
-            new ExpressionVirtualColumn(
+            makeExpressionVirtualColumn("concat(substring(countryIsoCode, 0, 1),'L')")
-                "virtual",
-                "concat(substring(countryIsoCode, 0, 1),'L')",
-                ValueType.STRING,
-                ExprMacroTable.nil()
-            )
         )
     );
 
@@ -1563,12 +1559,7 @@ public class HashJoinSegmentStorageAdapterTest extends BaseHashJoinSegmentStorag
 
     VirtualColumns virtualColumns = VirtualColumns.create(
         Collections.singletonList(
-            new ExpressionVirtualColumn(
+            makeExpressionVirtualColumn("concat(substring(countryIsoCode, 0, 1),'L')")
-                "virtual",
-                "concat(substring(countryIsoCode, 0, 1),'L')",
-                ValueType.STRING,
-                ExprMacroTable.nil()
-            )
         )
     );
 
@@ -2028,4 +2019,40 @@ public class HashJoinSegmentStorageAdapterTest extends BaseHashJoinSegmentStorag
     );
   }
 
+  @Test
+  public void test_determineBaseColumnsWithPreAndPostJoinVirtualColumns()
+  {
+    List<JoinableClause> joinableClauses = ImmutableList.of(factToCountryOnIsoCode(JoinType.LEFT));
+    JoinFilterPreAnalysis analysis = makeDefaultConfigPreAnalysis(null, joinableClauses, VirtualColumns.EMPTY);
+    HashJoinSegmentStorageAdapter adapter = new HashJoinSegmentStorageAdapter(
+        factSegment.asStorageAdapter(),
+        joinableClauses,
+        analysis
+    );
+    List<VirtualColumn> expectedPreJoin = ImmutableList.of(
+        makeExpressionVirtualColumn("concat(countryIsoCode,'L')", "v0"),
+        makeExpressionVirtualColumn("concat(countryIsoCode, countryNumber)", "v1"),
+        makeExpressionVirtualColumn("channel_uniques - 1", "v2"),
+        makeExpressionVirtualColumn("channel_uniques - __time", "v3")
+    );
+
+    List<VirtualColumn> expectedPostJoin = ImmutableList.of(
+        makeExpressionVirtualColumn("concat(countryIsoCode, dummyColumn)", "v4"),
+        makeExpressionVirtualColumn("dummyMetric - __time", "v5")
+    );
+    List<VirtualColumn> actualPreJoin = new ArrayList<>();
+    List<VirtualColumn> actualPostJoin = new ArrayList<>();
+    List<VirtualColumn> allVirtualColumns = new ArrayList<>();
+    allVirtualColumns.addAll(expectedPreJoin);
+    allVirtualColumns.addAll(expectedPostJoin);
+    adapter.determineBaseColumnsWithPreAndPostJoinVirtualColumns(
+        VirtualColumns.create(allVirtualColumns),
+        actualPreJoin,
+        actualPostJoin
+    );
+
+    Assert.assertEquals(expectedPreJoin, actualPreJoin);
+    Assert.assertEquals(expectedPostJoin, actualPostJoin);
+  }
+
 }