Make IDSET and IN_ID_SET work in the multi-stage engine

pinot-common/src/main/java/org/apache/pinot/common/function/TransformFunctionType.java

@@ -176,7 +176,10 @@ public enum TransformFunctionType {
   MAP_VALUE("mapValue",
       ReturnTypes.cascade(opBinding -> positionalComponentType(opBinding, 2), SqlTypeTransforms.FORCE_NULLABLE),
       OperandTypes.family(List.of(SqlTypeFamily.ARRAY, SqlTypeFamily.ANY, SqlTypeFamily.ARRAY))),
-  IN_ID_SET("inIdSet", ReturnTypes.BOOLEAN, OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)),
+  // Nullable: both the transform and the scalar function return NULL for a NULL value, unless the IdSet is empty. The
+  // multi-stage planner also uses this operator for the inIdSet scalar function. This operator is deterministic and
+  // the scalar function's is not, so this one lets the planner push IN_ID_SET below joins (see IdSetFunctions).
+  IN_ID_SET("inIdSet", ReturnTypes.BOOLEAN_NULLABLE, OperandTypes.family(SqlTypeFamily.ANY, SqlTypeFamily.CHARACTER)),
   LOOKUP("lookUp"),
   GROOVY("groovy"),
   SCALAR("scalar"),

pinot-core/src/main/java/org/apache/pinot/core/function/scalar/IdSetFunctions.java

@@ -0,0 +1,169 @@
+/**
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pinot.core.function.scalar;
+
+import com.google.common.annotations.VisibleForTesting;
+import java.io.IOException;
+import java.math.BigDecimal;
+import java.sql.Timestamp;
+import java.util.UUID;
+import javax.annotation.Nullable;
+import org.apache.pinot.core.query.utils.idset.IdSet;
+import org.apache.pinot.core.query.utils.idset.IdSets;
+import org.apache.pinot.spi.annotations.ScalarFunction;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.ByteArray;
+import org.apache.pinot.spi.utils.UuidUtils;
+
+
+/// Scalar form of the `IN_ID_SET` transform function, for queries that cannot run the transform function:
+/// - In the multi-stage engine, wherever the predicate is not pushed into a leaf stage, e.g. a filter over a join,
+///   `HAVING` or a join condition. Leaf stages keep using the `InIdSetTransformFunction`, because transform functions
+///   take precedence over scalar functions.
+/// - In the single-stage engine, after aggregation, e.g. in `HAVING`. Aggregations such as `MAX` return DOUBLE there,
+///   so cast them to the type the IdSet was built from.
+///
+/// The method is not static, so each call site gets its own instance (see `FunctionInvoker`), which caches the last
+/// deserialized IdSet:
+/// - A literal IdSet reaches every row as the same String instance, so it is deserialized once per call site.
+/// - An IdSet read from rows (e.g. a joined `IDSET` subquery result) arrives as a new but equal String instance with
+///   every data block, which costs one String comparison per block.
+/// - A call site whose IdSet changes from row to row (e.g. a correlated `IDSET` subquery, with one IdSet per group)
+///   deserializes the IdSet on every change, so it only suits small IdSets.
+///
+/// Not thread-safe: an instance serves one call site, which one thread evaluates at a time.
+public class IdSetFunctions {
+  private String _serializedIdSet;
+  private IdSet _idSet;
+  // The IdSet and the value class whose combination passed the last value type validation
+  private IdSet _validatedIdSet;
+  private Class<?> _validatedValueClass;
+
+  /// Returns whether the IdSet contains the value. As with an `IN` subquery, the result is `false` for an empty or
+  /// `null` IdSet, even when the value is `null`, and `null` for a `null` value otherwise. `IDSET` returns a `null`
+  /// IdSet over zero rows when null handling is enabled.
+  ///
+  /// Values arrive as their external Java types. BOOLEAN, TIMESTAMP and UUID values are looked up through their
+  /// stored INT, LONG and BYTES values, as the transform function does.
+  ///
+  /// `isDeterministic = false` keeps all-literal calls out of compile-time evaluation, which passes literals as Java
+  /// types the lookup cannot map to a value type (e.g. BigDecimal for an INT literal). The runtime evaluates them
+  /// instead, as before this function existed. The flag also marks the function VOLATILE, so ingestion transforms
+  /// cannot use it.
+  ///
+  /// The flag does not stop the multi-stage planner from pushing `IN_ID_SET` below a join into a leaf stage, because
+  /// the planner uses the operator of the `IN_ID_SET` transform function, which is deterministic. The operator table
+  /// registers transform functions before scalar functions, and skips this function because its name is taken. If
+  /// the transform function's operator is removed, this function's volatile operator takes its place, and
+  /// `PinotRuleUtils.isRelocatable` then keeps `IN_ID_SET` above joins. Tests in `JoinPlans.json` and
+  /// `PhysicalOptimizerPlans.json` pin the push-down.
+  @Nullable
+  @ScalarFunction(nullableParameters = true, isDeterministic = false)
+  public Boolean inIdSet(@Nullable Object value, @Nullable String serializedIdSet) {
+    if (serializedIdSet == null) {
+      return false;
+    }
+    IdSet idSet = getIdSet(serializedIdSet);
+    if (value == null) {
+      return idSet.getType() == IdSet.Type.EMPTY ? false : null;
+    }
+    if (idSet != _validatedIdSet || value.getClass() != _validatedValueClass) {
+      IdSets.validateValueType(idSet, getStoredType(value));
+      _validatedIdSet = idSet;
+      _validatedValueClass = value.getClass();
+    }
+    if (value instanceof Integer) {
+      return idSet.contains(((Integer) value).intValue());
+    }
+    if (value instanceof Long) {
+      return idSet.contains(((Long) value).longValue());
+    }
+    if (value instanceof Float) {
+      return idSet.contains(((Float) value).floatValue());
+    }
+    if (value instanceof Double) {
+      return idSet.contains(((Double) value).doubleValue());
+    }
+    if (value instanceof String) {
+      return idSet.contains((String) value);
+    }
+    if (value instanceof byte[]) {
+      return idSet.contains((byte[]) value);
+    }
+    if (value instanceof Boolean) {
+      return idSet.contains((Boolean) value ? 1 : 0);
+    }
+    if (value instanceof Timestamp) {
+      return idSet.contains(((Timestamp) value).getTime());
+    }
+    if (value instanceof UUID) {
+      return idSet.contains(UuidUtils.toBytes((UUID) value));
+    }
+    if (value instanceof ByteArray) {
+      return idSet.contains(((ByteArray) value).getBytes());
+    }
+    throw new IllegalStateException("Unsupported value class: " + value.getClass().getName());
+  }
+
+  @VisibleForTesting
+  IdSet getIdSet(String serializedIdSet) {
+    if (serializedIdSet != _serializedIdSet) {
+      if (!serializedIdSet.equals(_serializedIdSet)) {
+        _idSet = deserialize(serializedIdSet);
+      }
+      // Keep the latest instance: the following rows of a data block share it, so they match by reference
+      _serializedIdSet = serializedIdSet;
+    }
+    return _idSet;
+  }
+
+  private static IdSet deserialize(String serializedIdSet) {
+    try {
+      return IdSets.fromBase64String(serializedIdSet);
+    } catch (IOException | RuntimeException e) {
+      throw new IllegalArgumentException("Caught exception while deserializing IdSet: " + e, e);
+    }
+  }
+
+  private static DataType getStoredType(Object value) {
+    if (value instanceof Integer || value instanceof Boolean) {
+      return DataType.INT;
+    }
+    if (value instanceof Long || value instanceof Timestamp) {
+      return DataType.LONG;
+    }
+    if (value instanceof Float) {
+      return DataType.FLOAT;
+    }
+    if (value instanceof Double) {
+      return DataType.DOUBLE;
+    }
+    if (value instanceof String) {
+      return DataType.STRING;
+    }
+    if (value instanceof byte[] || value instanceof UUID || value instanceof ByteArray) {
+      return DataType.BYTES;
+    }
+    if (value instanceof BigDecimal) {
+      return DataType.BIG_DECIMAL;
+    }
+    throw new IllegalArgumentException(
+        "Cannot look up values of class " + value.getClass().getName() + " in an IdSet");
+  }
+}

pinot-core/src/main/java/org/apache/pinot/core/operator/transform/function/InIdSetTransformFunction.java

@@ -22,13 +22,15 @@
 import java.io.IOException;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 import org.apache.pinot.common.function.TransformFunctionType;
 import org.apache.pinot.core.operator.ColumnContext;
 import org.apache.pinot.core.operator.blocks.ValueBlock;
 import org.apache.pinot.core.operator.transform.TransformResultMetadata;
 import org.apache.pinot.core.query.utils.idset.IdSet;
 import org.apache.pinot.core.query.utils.idset.IdSets;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.roaringbitmap.RoaringBitmap;
 
 
 /// The IN_ID_SET transform function takes 2 arguments:
@@ -61,16 +63,24 @@ public void init(List<TransformFunction> arguments, Map<String, ColumnContext> c
     _transformFunction = arguments.get(0);
     try {
       _idSet = IdSets.fromBase64String(((LiteralTransformFunction) arguments.get(1)).getStringLiteral());
-    } catch (IOException e) {
-      throw new IllegalArgumentException("Caught exception while deserializing IdSet", e);
+    } catch (IOException | RuntimeException e) {
+      throw new IllegalArgumentException("Caught exception while deserializing IdSet: " + e, e);
     }
+    IdSets.validateValueType(_idSet, _transformFunction.getResultMetadata().getDataType().getStoredType());
   }
 
   @Override
   public TransformResultMetadata getResultMetadata() {
     return BOOLEAN_SV_NO_DICTIONARY_METADATA;
   }
 
+  /// As with an `IN` subquery, a NULL value is in no empty IdSet, so the result is never NULL for an empty IdSet.
+  @Nullable
+  @Override
+  public RoaringBitmap getNullBitmap(ValueBlock valueBlock) {
+    return _idSet.getType() == IdSet.Type.EMPTY ? null : super.getNullBitmap(valueBlock);
+  }
+
   @Override
   public int[] transformToIntValuesSV(ValueBlock valueBlock) {
     int length = valueBlock.getNumDocs();

pinot-core/src/main/java/org/apache/pinot/core/query/utils/idset/BloomFilterIdSet.java

@@ -18,14 +18,14 @@
  */
 package org.apache.pinot.core.query.utils.idset;
 
-import com.google.common.base.Preconditions;
 import com.google.common.hash.BloomFilter;
 import com.google.common.hash.Funnel;
 import com.google.common.hash.Funnels;
-import java.io.ByteArrayInputStream;
+import java.io.BufferedInputStream;
 import java.io.ByteArrayOutputStream;
 import java.io.IOException;
 import java.nio.ByteBuffer;
+import java.util.List;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 
 
@@ -36,12 +36,16 @@ public class BloomFilterIdSet implements IdSet {
 
   private enum FunnelType {
     // DO NOT change the ids as the ser/de relies on them
-    INT((byte) 0), LONG((byte) 1), STRING((byte) 2), BYTES((byte) 3);
+    INT((byte) 0, DataType.INT, DataType.FLOAT), LONG((byte) 1, DataType.LONG, DataType.DOUBLE),
+    STRING((byte) 2, DataType.STRING), BYTES((byte) 3, DataType.BYTES);
 
     private final byte _id;
+    // Stored types whose values hash through this funnel: FLOAT and DOUBLE values hash as their raw bits
+    private final List<DataType> _valueTypes;
 
-    FunnelType(byte id) {
+    FunnelType(byte id, DataType... valueTypes) {
       _id = id;
+      _valueTypes = List.of(valueTypes);
     }
 
     public byte getId() {
@@ -93,6 +97,12 @@ BloomFilter getBloomFilter() {
     return _bloomFilter;
   }
 
+  /// Returns the stored types whose values can be looked up in this IdSet. The serialized form records only the
+  /// funnel, so an IdSet built from INT values also accepts FLOAT values (and LONG also accepts DOUBLE).
+  List<DataType> getValueTypes() {
+    return _funnelType._valueTypes;
+  }
+
   @Override
   public Type getType() {
     return Type.BLOOM_FILTER;
@@ -174,13 +184,11 @@ public byte[] toBytes()
     return byteArrayOutputStream.toByteArray();
   }
 
-  /// Deserializes the BloomFilterIdSet from a ByteBuffer.
+  /// Deserializes the BloomFilterIdSet from a ByteBuffer, which may be read-only or direct.
   ///
   /// NOTE: The ByteBuffer does not include the IdSet.Type byte.
   static BloomFilterIdSet fromByteBuffer(ByteBuffer byteBuffer)
       throws IOException {
-    Preconditions.checkArgument(byteBuffer.hasArray(),
-        "Cannot deserialize BloomFilter from ByteBuffer not backed by an accessible byte array");
     // Count the IdSet.Type byte
     int serializedSizeInBytes = 1 + byteBuffer.remaining();
     byte funnelTypeId = byteBuffer.get();
@@ -204,12 +212,11 @@ static BloomFilterIdSet fromByteBuffer(ByteBuffer byteBuffer)
         funnel = Funnels.byteArrayFunnel();
         break;
       default:
-        throw new IllegalStateException();
+        throw new IllegalStateException("Unsupported BloomFilter funnel type id: " + funnelTypeId);
     }
     // NOTE: No need to close the stream.
-    BloomFilter bloomFilter = BloomFilter.readFrom(
-        new ByteArrayInputStream(byteBuffer.array(), byteBuffer.arrayOffset() + byteBuffer.position(),
-            byteBuffer.remaining()), funnel);
+    // Guava reads the filter 8 bytes at a time, which the Pinot buffer stream copies byte by byte, so buffer the reads.
+    BloomFilter bloomFilter = BloomFilter.readFrom(new BufferedInputStream(IdSets.toInputStream(byteBuffer)), funnel);
     return new BloomFilterIdSet(funnelType, bloomFilter, serializedSizeInBytes);
   }
 

pinot-core/src/main/java/org/apache/pinot/core/query/utils/idset/IdSets.java

@@ -21,6 +21,11 @@
 import java.io.IOException;
 import java.nio.ByteBuffer;
 import java.util.Base64;
+import java.util.List;
+import java.util.stream.Collectors;
+import org.apache.pinot.segment.spi.memory.DataBufferPinotInputStream;
+import org.apache.pinot.segment.spi.memory.PinotByteBuffer;
+import org.apache.pinot.segment.spi.memory.PinotInputStream;
 import org.apache.pinot.spi.data.FieldSpec.DataType;
 import org.roaringbitmap.PeekableIntIterator;
 import org.roaringbitmap.longlong.LongIterator;
@@ -38,6 +43,10 @@ private IdSets() {
   public static final int DEFAULT_EXPECTED_INSERTIONS = 5_000_000;
   public static final double DEFAULT_FPP = 0.03;
 
+  // Stored types whose values can be looked up in an IdSet
+  private static final List<DataType> LOOKUP_VALUE_TYPES =
+      List.of(DataType.INT, DataType.LONG, DataType.FLOAT, DataType.DOUBLE, DataType.STRING, DataType.BYTES);
+
   /// Creates an IdSet.
   ///
   /// @param dataType Data type of the ids
@@ -78,6 +87,40 @@ public static EmptyIdSet emptyIdSet() {
     return EmptyIdSet.INSTANCE;
   }
 
+  /// Validates that values of the given stored type can be looked up in the given IdSet. Rejects e.g. LONG values
+  /// in an IdSet built from INT values, where every lookup would otherwise throw or hash to the wrong bits. An empty
+  /// IdSet accepts every lookup type, as it records no value type.
+  ///
+  /// @throws IllegalArgumentException if values of the stored type cannot be looked up in the IdSet
+  public static void validateValueType(IdSet idSet, DataType storedType) {
+    if (!LOOKUP_VALUE_TYPES.contains(storedType)) {
+      throw new IllegalArgumentException(
+          String.format("Cannot look up %s values in an IdSet, supported value types: %s", storedType,
+              LOOKUP_VALUE_TYPES));
+    }
+    List<DataType> valueTypes;
+    switch (idSet.getType()) {
+      case EMPTY:
+        return;
+      case ROARING_BITMAP:
+        valueTypes = List.of(DataType.INT);
+        break;
+      case ROARING_64_NAVIGABLE_MAP:
+        valueTypes = List.of(DataType.LONG);
+        break;
+      case BLOOM_FILTER:
+        valueTypes = ((BloomFilterIdSet) idSet).getValueTypes();
+        break;
+      default:
+        throw new IllegalStateException("Unsupported IdSet type: " + idSet.getType());
+    }
+    if (!valueTypes.contains(storedType)) {
+      throw new IllegalArgumentException(
+          String.format("Cannot look up %s values in an IdSet built from %s values", storedType,
+              valueTypes.stream().map(DataType::name).collect(Collectors.joining(" or "))));
+    }
+  }
+
   /// Merges 2 IdSets, converts to BloomFilterIdSet if the size exceeds the size threshold.
   ///
   /// @param idSet1 First IdSet to merge
@@ -187,7 +230,7 @@ public static IdSet fromBytes(byte[] bytes)
     return fromByteBuffer(ByteBuffer.wrap(bytes));
   }
 
-  /// Deserializes the IdSet from a ByteBuffer.
+  /// Deserializes the IdSet from a ByteBuffer, which may be read-only or direct (e.g. a multi-stage mailbox payload).
   public static IdSet fromByteBuffer(ByteBuffer byteBuffer)
       throws IOException {
     byte typeId = byteBuffer.get();
@@ -201,7 +244,7 @@ public static IdSet fromByteBuffer(ByteBuffer byteBuffer)
       case 3:
         return BloomFilterIdSet.fromByteBuffer(byteBuffer);
       default:
-        throw new IllegalStateException();
+        throw new IllegalStateException("Unsupported IdSet type id: " + typeId);
     }
   }
 
@@ -212,4 +255,10 @@ public static IdSet fromBase64String(String base64String)
       throws IOException {
     return fromBytes(Base64.getDecoder().decode(base64String));
   }
+
+  /// Returns a [PinotInputStream] over the remaining bytes of the given buffer, without moving its position. It reads
+  /// heap, direct and read-only buffers the same way, and its `DataInput` methods read big-endian values.
+  static PinotInputStream toInputStream(ByteBuffer byteBuffer) {
+    return new DataBufferPinotInputStream(PinotByteBuffer.slice(byteBuffer));
+  }
 }