private LazyBitmap getMatchingFlattenedDocIds()

in pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/json/MutableJsonIndexImpl.java [320:509]
154 lines of code
45 McCabe index (conditional complexity)

  private LazyBitmap getMatchingFlattenedDocIds(Predicate predicate) {
    ExpressionContext lhs = predicate.getLhs();
    Preconditions.checkArgument(lhs.getType() == ExpressionContext.Type.IDENTIFIER,
        "Left-hand side of the predicate must be an identifier, got: %s (%s). Put double quotes around the identifier"
            + " if needed.", lhs, lhs.getType());
    String key = lhs.getIdentifier();

    // Support 2 formats:
    // - JSONPath format (e.g. "$.a[1].b"='abc', "$[0]"=1, "$"='abc')
    // - Legacy format (e.g. "a[1].b"='abc')
    if (key.charAt(0) == '$') {
      key = key.substring(1);
    } else {
      key = JsonUtils.KEY_SEPARATOR + key;
    }
    Pair<String, LazyBitmap> pair = getKeyAndFlattenedDocIds(key);
    key = pair.getLeft();
    LazyBitmap matchingDocIds = pair.getRight();
    if (matchingDocIds != null && matchingDocIds.isEmpty()) {
      return LazyBitmap.EMPTY_BITMAP;
    }

    Predicate.Type predicateType = predicate.getType();
    switch (predicateType) {
      case EQ: {
        String value = ((EqPredicate) predicate).getValue();
        String keyValuePair = key + JsonIndexCreator.KEY_VALUE_SEPARATOR + value;
        RoaringBitmap result = _postingListMap.get(keyValuePair);
        return filter(result, matchingDocIds);
      }

      case NOT_EQ: {
        String notEqualValue = ((NotEqPredicate) predicate).getValue();
        LazyBitmap result = null;

        RoaringBitmap allDocIds = _postingListMap.get(key);
        if (allDocIds != null && !allDocIds.isEmpty()) {
          result = new LazyBitmap(allDocIds);

          RoaringBitmap notEqualDocIds =
              _postingListMap.get(key + JsonIndexCreator.KEY_VALUE_SEPARATOR + notEqualValue);

          if (notEqualDocIds != null && !notEqualDocIds.isEmpty()) {
            result = result.andNot(notEqualDocIds);
          }
        }

        return filter(result, matchingDocIds);
      }

      case IN: {
        List<String> values = ((InPredicate) predicate).getValues();
        LazyBitmap result = null;

        StringBuilder buffer = new StringBuilder(key);
        buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR);
        int pos = buffer.length();

        for (String value : values) {
          buffer.setLength(pos);
          buffer.append(value);
          String keyValue = buffer.toString();

          RoaringBitmap docIds = _postingListMap.get(keyValue);

          if (docIds != null && !docIds.isEmpty()) {
            if (result == null) {
              result = new LazyBitmap(docIds);
            } else {
              result = result.or(docIds);
            }
          }
        }

        return filter(result, matchingDocIds);
      }

      case NOT_IN: {
        List<String> notInValues = ((NotInPredicate) predicate).getValues();
        LazyBitmap result = null;

        RoaringBitmap allDocIds = _postingListMap.get(key);
        if (allDocIds != null && !allDocIds.isEmpty()) {
          result = new LazyBitmap(allDocIds);

          StringBuilder buffer = new StringBuilder(key);
          buffer.append(JsonIndexCreator.KEY_VALUE_SEPARATOR);
          int pos = buffer.length();

          for (String notInValue : notInValues) {
            buffer.setLength(pos);
            buffer.append(notInValue);
            String keyValuePair = buffer.toString();

            RoaringBitmap docIds = _postingListMap.get(keyValuePair);
            if (docIds != null && !docIds.isEmpty()) {
              result = result.andNot(docIds);
            }
          }
        }

        return filter(result, matchingDocIds);
      }

      case IS_NOT_NULL:
      case IS_NULL: {
        RoaringBitmap result = _postingListMap.get(key);
        return filter(result, matchingDocIds);
      }

      case REGEXP_LIKE: {
        Map<String, RoaringBitmap> subMap = getMatchingKeysMap(key);
        if (subMap.isEmpty()) {
          return LazyBitmap.EMPTY_BITMAP;
        }

        Pattern pattern = ((RegexpLikePredicate) predicate).getPattern();
        Matcher matcher = pattern.matcher("");
        LazyBitmap result = null;
        StringBuilder value = new StringBuilder();

        for (Map.Entry<String, RoaringBitmap> entry : subMap.entrySet()) {
          String keyValue = entry.getKey();
          value.setLength(0);
          value.append(keyValue, key.length() + 1, keyValue.length());

          if (!matcher.reset(value).matches()) {
            continue;
          }

          RoaringBitmap docIds = entry.getValue();
          if (docIds != null && !docIds.isEmpty()) {
            if (result == null) {
              result = new LazyBitmap(docIds);
            } else {
              result = result.or(docIds);
            }
          }
        }

        return filter(result, matchingDocIds);
      }

      case RANGE: {
        Map<String, RoaringBitmap> subMap = getMatchingKeysMap(key);
        if (subMap.isEmpty()) {
          return LazyBitmap.EMPTY_BITMAP;
        }

        LazyBitmap result = null;
        RangePredicate rangePredicate = (RangePredicate) predicate;
        FieldSpec.DataType rangeDataType = rangePredicate.getRangeDataType();
        // Simplify to only support numeric and string types
        if (rangeDataType.isNumeric()) {
          rangeDataType = FieldSpec.DataType.DOUBLE;
        } else {
          rangeDataType = FieldSpec.DataType.STRING;
        }

        boolean lowerUnbounded = rangePredicate.getLowerBound().equals(RangePredicate.UNBOUNDED);
        boolean upperUnbounded = rangePredicate.getUpperBound().equals(RangePredicate.UNBOUNDED);
        boolean lowerInclusive = lowerUnbounded || rangePredicate.isLowerInclusive();
        boolean upperInclusive = upperUnbounded || rangePredicate.isUpperInclusive();
        Object lowerBound = lowerUnbounded ? null : rangeDataType.convert(rangePredicate.getLowerBound());
        Object upperBound = upperUnbounded ? null : rangeDataType.convert(rangePredicate.getUpperBound());

        for (Map.Entry<String, RoaringBitmap> entry : subMap.entrySet()) {
          Object valueObj = rangeDataType.convert(entry.getKey().substring(key.length() + 1));
          boolean lowerCompareResult =
              lowerUnbounded || (lowerInclusive ? rangeDataType.compare(valueObj, lowerBound) >= 0
                  : rangeDataType.compare(valueObj, lowerBound) > 0);
          boolean upperCompareResult =
              upperUnbounded || (upperInclusive ? rangeDataType.compare(valueObj, upperBound) <= 0
                  : rangeDataType.compare(valueObj, upperBound) < 0);
          if (lowerCompareResult && upperCompareResult) {
            if (result == null) {
              result = new LazyBitmap(entry.getValue());
            } else {
              result = result.or(entry.getValue());
            }
          }
        }

        return filter(result, matchingDocIds);
      }

      default:
        throw new IllegalStateException("Unsupported json_match predicate type: " + predicate);
    }
  }