QueryableIndexColumnSelectorFactory: Double-check cached column class. (#11957)

Important because an earlier call to getCachedColumn may have been
done with a different class, leading to a ClassCastException on the
second call. In the prior code, this could happen if a complex column
had makeDimensionSelector called on it after makeColumnValueSelector had
already been called.

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

@@ -172,7 +172,7 @@ public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactor
   @SuppressWarnings("unchecked")
   private <T extends BaseColumn> T getCachedColumn(final String columnName, final Class<T> clazz)
   {
-    return (T) columnCache.computeIfAbsent(
+    final BaseColumn cachedColumn = columnCache.computeIfAbsent(
         columnName,
         name -> {
           ColumnHolder holder = index.getColumnHolder(name);
@@ -185,6 +185,12 @@ public class QueryableIndexColumnSelectorFactory implements ColumnSelectorFactor
           }
         }
     );
+
+    if (cachedColumn != null && clazz.isAssignableFrom(cachedColumn.getClass())) {
+      return (T) cachedColumn;
+    } else {
+      return null;
+    }
   }
 
   @Override

processing/src/test/java/org/apache/druid/segment/QueryableIndexStorageAdapterTest.java

@@ -19,6 +19,7 @@
 
 package org.apache.druid.segment;
 
+import org.apache.druid.hll.HyperLogLogCollector;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.java.util.common.granularity.Granularities;
@@ -30,6 +31,8 @@ import org.apache.druid.query.dimension.DefaultDimensionSpec;
 import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
 import org.apache.druid.segment.vector.VectorCursor;
 import org.apache.druid.testing.InitializedNullHandlingTest;
+import org.hamcrest.CoreMatchers;
+import org.hamcrest.MatcherAssert;
 import org.junit.After;
 import org.junit.Assert;
 import org.junit.Before;
@@ -237,4 +240,69 @@ public class QueryableIndexStorageAdapterTest
       Assert.assertTrue(partialNullSelector.supportsLookupNameUtf8());
     }
   }
+
+  public static class ManySelectorsOneColumnTest extends InitializedNullHandlingTest
+  {
+    private Cursor cursor;
+    private ColumnSelectorFactory columnSelectorFactory;
+    private final Closer closer = Closer.create();
+
+    @Before
+    public void setUp()
+    {
+      final QueryableIndex index = TestIndex.getMMappedTestIndex();
+      final QueryableIndexStorageAdapter adapter = new QueryableIndexStorageAdapter(index);
+      final Sequence<Cursor> cursors = adapter.makeCursors(
+          null,
+          Intervals.ETERNITY,
+          VirtualColumns.EMPTY,
+          Granularities.ALL,
+          false,
+          null
+      );
+      final Yielder<Cursor> cursorYielder = Yielders.each(cursors);
+      cursor = cursorYielder.get();
+      columnSelectorFactory = cursor.getColumnSelectorFactory();
+      closer.register(cursorYielder);
+    }
+
+    @After
+    public void testDown() throws IOException
+    {
+      closer.close();
+    }
+
+    @Test
+    public void testTwoSelectorsOneComplexColumn()
+    {
+      final ColumnValueSelector<?> valueSelector = columnSelectorFactory.makeColumnValueSelector("quality_uniques");
+      MatcherAssert.assertThat(valueSelector.getObject(), CoreMatchers.instanceOf(HyperLogLogCollector.class));
+
+      final DimensionSelector dimensionSelector =
+          columnSelectorFactory.makeDimensionSelector(DefaultDimensionSpec.of("quality_uniques"));
+      Assert.assertNull(dimensionSelector.getObject());
+    }
+
+    @Test
+    public void testTwoSelectorsOneNumericColumn()
+    {
+      final ColumnValueSelector<?> valueSelector = columnSelectorFactory.makeColumnValueSelector("index");
+      MatcherAssert.assertThat(valueSelector.getObject(), CoreMatchers.instanceOf(Double.class));
+
+      final DimensionSelector dimensionSelector =
+          columnSelectorFactory.makeDimensionSelector(DefaultDimensionSpec.of("index"));
+      Assert.assertEquals("100.0", dimensionSelector.getObject());
+    }
+
+    @Test
+    public void testTwoSelectorsOneStringColumn()
+    {
+      final ColumnValueSelector<?> valueSelector = columnSelectorFactory.makeColumnValueSelector("market");
+      MatcherAssert.assertThat(valueSelector.getObject(), CoreMatchers.instanceOf(String.class));
+
+      final DimensionSelector dimensionSelector =
+          columnSelectorFactory.makeDimensionSelector(DefaultDimensionSpec.of("market"));
+      MatcherAssert.assertThat(dimensionSelector.getObject(), CoreMatchers.instanceOf(String.class));
+    }
+  }
 }