fix bug when front-coded index has only the null value (#13309)

processing/src/main/java/org/apache/druid/segment/data/FixedIndexed.java

@@ -22,7 +22,6 @@ package org.apache.druid.segment.data;
 import com.google.common.base.Preconditions;
 import com.google.common.base.Supplier;
 import org.apache.druid.common.config.NullHandling;
-import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import org.apache.druid.segment.column.TypeStrategy;
 
@@ -119,12 +118,7 @@ public class FixedIndexed<T> implements Indexed<T>
   @Override
   public T get(int index)
   {
-    if (index < 0) {
-      throw new IAE("Index[%s] < 0", index);
-    }
-    if (index >= size) {
-      throw new IAE("Index[%d] >= size[%d]", index, size);
-    }
+    Indexed.checkIndex(index, size);
     if (hasNull) {
       if (index == 0) {
         return null;

processing/src/main/java/org/apache/druid/segment/data/FrontCodedIndexed.java

@@ -28,6 +28,7 @@ import org.apache.druid.query.monomorphicprocessing.RuntimeShapeInspector;
 import javax.annotation.Nullable;
 import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.NoSuchElementException;
 
@@ -146,6 +147,7 @@ public final class FrontCodedIndexed implements Indexed<ByteBuffer>
     if (hasNull && index == 0) {
       return null;
     }
+    Indexed.checkIndex(index, adjustedNumValues);
 
     // due to vbyte encoding, the null value is not actually stored in the bucket (no negative values), so we adjust
     // the index
@@ -266,6 +268,9 @@ public final class FrontCodedIndexed implements Indexed<ByteBuffer>
   @Override
   public Iterator<ByteBuffer> iterator()
   {
+    if (hasNull && adjustedNumValues == 1) {
+      return Collections.<ByteBuffer>singletonList(null).iterator();
+    }
     ByteBuffer copy = buffer.asReadOnlyBuffer().order(buffer.order());
     copy.position(bucketsPosition);
     final ByteBuffer[] firstBucket = readBucket(copy, numBuckets > 1 ? bucketSize : lastBucketNumValues);

processing/src/main/java/org/apache/druid/segment/data/FrontCodedIndexedWriter.java

@@ -224,6 +224,9 @@ public class FrontCodedIndexedWriter implements DictionaryWriter<byte[]>
 
   private void flush() throws IOException
   {
+    if (numWritten == 0) {
+      return;
+    }
     int remainder = numWritten % bucketSize;
     resetScratch();
     int written;

processing/src/main/java/org/apache/druid/segment/data/Indexed.java

@@ -20,6 +20,7 @@
 package org.apache.druid.segment.data;
 
 import org.apache.druid.guice.annotations.PublicApi;
+import org.apache.druid.java.util.common.IAE;
 import org.apache.druid.query.monomorphicprocessing.CalledFromHotLoop;
 import org.apache.druid.query.monomorphicprocessing.HotLoopCallee;
 
@@ -67,4 +68,23 @@ public interface Indexed<T> extends Iterable<T>, HotLoopCallee
   {
     return false;
   }
+
+  /**
+   * Checks  if {@code index} is between 0 and {@code size}. Similar to Preconditions.checkElementIndex() except this
+   * method throws {@link IAE} with custom error message.
+   * <p>
+   * Used here to get existing behavior(same error message and exception) of V1 {@link GenericIndexed}.
+   *
+   * @param index identifying an element of an {@link Indexed}
+   * @param size size of the {@link Indexed}
+   */
+  static void checkIndex(int index, int size)
+  {
+    if (index < 0) {
+      throw new IAE("Index[%s] < 0", index);
+    }
+    if (index >= size) {
+      throw new IAE("Index[%d] >= size[%d]", index, size);
+    }
+  }
 }

processing/src/test/java/org/apache/druid/segment/data/FrontCodedIndexedTest.java

@@ -34,6 +34,7 @@ import java.nio.ByteBuffer;
 import java.nio.ByteOrder;
 import java.nio.channels.WritableByteChannel;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.Iterator;
 import java.util.List;
 import java.util.TreeSet;
@@ -263,6 +264,32 @@ public class FrontCodedIndexedTest extends InitializedNullHandlingTest
     Assert.assertEquals(newListIterator.hasNext(), utf8Iterator.hasNext());
   }
 
+  @Test
+  public void testFrontCodedOnlyNull() throws IOException
+  {
+    ByteBuffer buffer = ByteBuffer.allocate(1 << 12).order(order);
+    List<String> theList = Collections.singletonList(null);
+    fillBuffer(buffer, theList, 4);
+
+    buffer.position(0);
+    FrontCodedIndexed codedUtf8Indexed = FrontCodedIndexed.read(
+        buffer,
+        buffer.order()
+    ).get();
+
+    Assert.assertNull(codedUtf8Indexed.get(0));
+    Assert.assertThrows(IllegalArgumentException.class, () -> codedUtf8Indexed.get(-1));
+    Assert.assertThrows(IllegalArgumentException.class, () -> codedUtf8Indexed.get(theList.size()));
+
+    Assert.assertEquals(0, codedUtf8Indexed.indexOf(null));
+    Assert.assertEquals(-2, codedUtf8Indexed.indexOf(StringUtils.toUtf8ByteBuffer("hello")));
+
+    Iterator<ByteBuffer> utf8Iterator = codedUtf8Indexed.iterator();
+    Assert.assertTrue(utf8Iterator.hasNext());
+    Assert.assertNull(utf8Iterator.next());
+    Assert.assertFalse(utf8Iterator.hasNext());
+  }
+
   private static long fillBuffer(ByteBuffer buffer, Iterable<String> sortedIterable, int bucketSize) throws IOException
   {
     Iterator<String> sortedStrings = sortedIterable.iterator();