com.tdunning.math.stats.TDigest Java Examples

@Test
public void testTDigest() {
    double compression = 100;
    double quantile = 0.5;

    PercentileCounter counter = new PercentileCounter(compression, quantile);
    TDigest tDigest = TDigest.createAvlTreeDigest(compression);

    Random random = new Random();
    int dataSize = 10000;
    List<Double> dataset = Lists.newArrayListWithCapacity(dataSize);
    for (int i = 0; i < dataSize; i++) {
        double d = random.nextDouble();
        counter.add(d);
        tDigest.add(d);
    }
    double actualResult = counter.getResultEstimate();

    Collections.sort(dataset);
    double expectedResult = tDigest.quantile(quantile);

    assertEquals(expectedResult, actualResult, 0);
}
 

@Test
public void testBasic() {
    int times = 1;
    int compression = 100;
    for (int t = 0; t < times; t++) {
        TDigest tDigest = TDigest.createAvlTreeDigest(compression);
        Random random = new Random();
        int dataSize = 10000;
        List<Double> dataset = Lists.newArrayListWithCapacity(dataSize);
        for (int i = 0; i < dataSize; i++) {
            double d = random.nextDouble();
            tDigest.add(d);
            dataset.add(d);
        }
        Collections.sort(dataset);

        double actualResult = tDigest.quantile(0.5);
        double expectedResult = MathUtil.findMedianInSortedList(dataset);
        assertEquals(expectedResult, actualResult, 0.01);
    }
}
 

@Test
public void testInnerSegmentGroupBy() {
  // For inner segment case, percentile does not affect the intermediate result
  AggregationGroupByOperator groupByOperator = getOperatorForQuery(getGroupByQuery(0));
  IntermediateResultsBlock resultsBlock = groupByOperator.nextBlock();
  AggregationGroupByResult groupByResult = resultsBlock.getAggregationGroupByResult();
  Assert.assertNotNull(groupByResult);
  Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = groupByResult.getGroupKeyIterator();
  while (groupKeyIterator.hasNext()) {
    GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
    DoubleList doubleList = (DoubleList) groupByResult.getResultForKey(groupKey, 0);
    Collections.sort(doubleList);
    assertTDigest((TDigest) groupByResult.getResultForKey(groupKey, 1), doubleList);
    assertTDigest((TDigest) groupByResult.getResultForKey(groupKey, 2), doubleList);
  }
}
 

private void compareQD(PrintWriter out, AbstractContinousDistribution gen, String tag, long scale) {
    for (double compression : new double[]{10, 20, 50, 100, 200, 500, 1000, 2000}) {
        QDigest qd = new QDigest(compression);
        TDigest dist = new MergingDigest(compression);
        double[] data = new double[100000];
        for (int i = 0; i < 100000; i++) {
            double x = gen.nextDouble();
            dist.add(x);
            qd.offer((long) (x * scale));
            data[i] = x;
        }
        dist.compress();
        Arrays.sort(data);

        for (double q : new double[]{1e-5, 1e-4, 0.001, 0.01, 0.1, 0.5, 0.9, 0.99, 0.999, 0.9999, 0.99999}) {
            double x1 = dist.quantile(q);
            double x2 = (double) qd.getQuantile(q) / scale;
            double e1 = Dist.cdf(x1, data) - q;
            double e2 = Dist.cdf(x2, data) - q;
            out.printf("%s,%.0f,%.8f,%.10g,%.10g,%d,%d\n", tag, compression, q, e1, e2, dist.smallByteSize(), QDigest.serialize(qd).length);
        }
    }
}
 

@Test
public void testTDigest() {
    double compression = 100;
    double quantile = 0.5;

    PercentileCounter counter = new PercentileCounter(compression, quantile);
    TDigest tDigest = TDigest.createAvlTreeDigest(compression);

    Random random = new Random();
    int dataSize = 10000;
    List<Double> dataset = Lists.newArrayListWithCapacity(dataSize);
    for (int i = 0; i < dataSize; i++) {
        double d = random.nextDouble();
        counter.add(d);
        tDigest.add(d);
    }
    double actualResult = counter.getResultEstimate();

    Collections.sort(dataset);
    double expectedResult = tDigest.quantile(quantile);

    assertEquals(expectedResult, actualResult, 0);
}
 

@Test
public void testBasic() {
    int times = 1;
    int compression = 100;
    for (int t = 0; t < times; t++) {
        TDigest tDigest = TDigest.createAvlTreeDigest(compression);
        Random random = new Random();
        int dataSize = 10000;
        List<Double> dataset = Lists.newArrayListWithCapacity(dataSize);
        for (int i = 0; i < dataSize; i++) {
            double d = random.nextDouble();
            tDigest.add(d);
            dataset.add(d);
        }
        Collections.sort(dataset);

        double actualResult = tDigest.quantile(0.5);
        double expectedResult = MathUtil.findMedianInSortedList(dataset);
        assertEquals(expectedResult, actualResult, 0.01);
    }
}
 

@Test
public void testTDigest() {
  for (int i = 0; i < NUM_ITERATIONS; i++) {
    TDigest expected = TDigest.createMergingDigest(PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION);
    int size = RANDOM.nextInt(100) + 1;
    for (int j = 0; j < size; j++) {
      expected.add(RANDOM.nextDouble());
    }

    byte[] bytes = ObjectSerDeUtils.serialize(expected);
    TDigest actual = ObjectSerDeUtils.deserialize(bytes, ObjectSerDeUtils.ObjectType.TDigest);

    for (int j = 0; j <= 100; j++) {
      assertEquals(actual.quantile(j / 100.0), expected.quantile(j / 100.0), 1e-5);
    }
  }
}
 

@Override
public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
    Map<ExpressionContext, BlockValSet> blockValSetMap) {
  BlockValSet blockValSet = blockValSetMap.get(_expression);
  if (blockValSet.getValueType() != DataType.BYTES) {
    double[] doubleValues = blockValSet.getDoubleValuesSV();
    for (int i = 0; i < length; i++) {
      getDefaultTDigest(groupByResultHolder, groupKeyArray[i]).add(doubleValues[i]);
    }
  } else {
    // Serialized TDigest
    byte[][] bytesValues = blockValSet.getBytesValuesSV();
    for (int i = 0; i < length; i++) {
      TDigest value = ObjectSerDeUtils.TDIGEST_SER_DE.deserialize(bytesValues[i]);
      int groupKey = groupKeyArray[i];
      TDigest tDigest = groupByResultHolder.getResult(groupKey);
      if (tDigest != null) {
        tDigest.add(value);
      } else {
        groupByResultHolder.setValueForKey(groupKey, value);
      }
    }
  }
}
 

@Override
public void serialize(TDigest td, JsonGenerator j, SerializerProvider sp) throws IOException, JsonProcessingException {
    try(ByteArrayOutputStream baos = new ByteArrayOutputStream(); ObjectOutputStream oos = new ObjectOutputStream(baos)){
        oos.writeObject(td);
        oos.close();
        byte[] bytes = baos.toByteArray();
        Base64 b = new Base64();
        String str = b.encodeAsString(bytes);
        j.writeStartObject();
        j.writeStringField("digest", str);
        j.writeEndObject();
    }
}
 

@Override
public TDigest deserialize(JsonParser jp, DeserializationContext d) throws IOException, JsonProcessingException {
    JsonNode node = (JsonNode)jp.getCodec().readTree(jp);
    String field = node.get("digest").asText();
    Base64 b = new Base64();
    byte[] bytes = b.decode(field);
    try(ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bytes))){
        return (TDigest) ois.readObject();
    } catch (Exception e){
        throw new RuntimeException("Error deserializing TDigest object from JSON", e);
    }
}
 

@Override
public TDigest deserialize(JsonParser jp, DeserializationContext d) throws IOException, JsonProcessingException {
    JsonNode node = (JsonNode)jp.getCodec().readTree(jp);
    String field = node.get("digest").asText();
    Base64 b = new Base64();
    byte[] bytes = b.decode(field);
    try(ObjectInputStream ois = new ObjectInputStream(new ByteArrayInputStream(bytes))){
        return (TDigest) ois.readObject();
    } catch (Exception e){
        throw new RuntimeException("Error deserializing TDigest object from JSON", e);
    }
}
 

@Override
public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
    Map<ExpressionContext, BlockValSet> blockValSetMap) {
  double[][] valuesArray = blockValSetMap.get(_expression).getDoubleValuesMV();
  TDigest tDigest = getDefaultTDigest(aggregationResultHolder);
  for (int i = 0; i < length; i++) {
    for (double value : valuesArray[i]) {
      tDigest.add(value);
    }
  }
}
 

@Override
public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
    Map<ExpressionContext, BlockValSet> blockValSetMap) {
  double[][] valuesArray = blockValSetMap.get(_expression).getDoubleValuesMV();
  for (int i = 0; i < length; i++) {
    TDigest tDigest = getDefaultTDigest(groupByResultHolder, groupKeyArray[i]);
    for (double value : valuesArray[i]) {
      tDigest.add(value);
    }
  }
}
 

@Override
public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
    Map<ExpressionContext, BlockValSet> blockValSetMap) {
  double[][] valuesArray = blockValSetMap.get(_expression).getDoubleValuesMV();
  for (int i = 0; i < length; i++) {
    double[] values = valuesArray[i];
    for (int groupKey : groupKeysArray[i]) {
      TDigest tDigest = getDefaultTDigest(groupByResultHolder, groupKey);
      for (double value : values) {
        tDigest.add(value);
      }
    }
  }
}
 

@Override
public TDigest extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
  TDigest tDigest = aggregationResultHolder.getResult();
  if (tDigest == null) {
    return TDigest.createMergingDigest(DEFAULT_TDIGEST_COMPRESSION);
  } else {
    return tDigest;
  }
}
 

/**
 * Returns the TDigest for the given group key if exists, or creates a new one with default compression.
 *
 * @param groupByResultHolder Result holder
 * @param groupKey Group key for which to return the TDigest
 * @return TDigest for the group key
 */
protected static TDigest getDefaultTDigest(GroupByResultHolder groupByResultHolder, int groupKey) {
  TDigest tDigest = groupByResultHolder.getResult(groupKey);
  if (tDigest == null) {
    tDigest = TDigest.createMergingDigest(DEFAULT_TDIGEST_COMPRESSION);
    groupByResultHolder.setValueForKey(groupKey, tDigest);
  }
  return tDigest;
}
 

private static TDigest getDefaultTDigest() {
  //return TDigest.createDigest(COMPRESSION);
  return new AVLTreeDigest(COMPRESSION);
  //return new TreeDigest(COMPRESSION);
  //return new ArrayDigest(4, COMPRESSION);
  //return new MergingDigest(COMPRESSION);    
}
 

@Override
public TDigest extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
  TDigest tDigest = groupByResultHolder.getResult(groupKey);
  if (tDigest == null) {
    return TDigest.createMergingDigest(DEFAULT_TDIGEST_COMPRESSION);
  } else {
    return tDigest;
  }
}
 

@Override
public TDigest merge(TDigest intermediateResult1, TDigest intermediateResult2) {
  if (intermediateResult1.size() == 0L) {
    return intermediateResult2;
  }
  if (intermediateResult2.size() == 0L) {
    return intermediateResult1;
  }
  intermediateResult1.add(intermediateResult2);
  return intermediateResult1;
}
 

/**
 * Returns the TDigest from the result holder or creates a new one with default compression if it does not exist.
 *
 * @param aggregationResultHolder Result holder
 * @return TDigest from the result holder
 */
protected static TDigest getDefaultTDigest(AggregationResultHolder aggregationResultHolder) {
  TDigest tDigest = aggregationResultHolder.getResult();
  if (tDigest == null) {
    tDigest = TDigest.createMergingDigest(DEFAULT_TDIGEST_COMPRESSION);
    aggregationResultHolder.setValue(tDigest);
  }
  return tDigest;
}
 

@Override
public TDigest getInitialAggregatedValue(Object rawValue) {
  TDigest initialValue;
  if (rawValue instanceof byte[]) {
    byte[] bytes = (byte[]) rawValue;
    initialValue = deserializeAggregatedValue(bytes);
    _maxByteSize = Math.max(_maxByteSize, bytes.length);
  } else {
    initialValue = TDigest.createMergingDigest(PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION);
    initialValue.add(((Number) rawValue).doubleValue());
    _maxByteSize = Math.max(_maxByteSize, initialValue.byteSize());
  }
  return initialValue;
}
 

@Override
public TDigest applyRawValue(TDigest value, Object rawValue) {
  if (rawValue instanceof byte[]) {
    value.add(deserializeAggregatedValue((byte[]) rawValue));
  } else {
    value.add(((Number) rawValue).doubleValue());
  }
  _maxByteSize = Math.max(_maxByteSize, value.byteSize());
  return value;
}
 

private void compareSQ(PrintWriter out, AbstractContinousDistribution gen, String tag) {
    double[] quantiles = {0.001, 0.01, 0.1, 0.2, 0.3, 0.5, 0.7, 0.8, 0.9, 0.99, 0.999};
    for (double compression : new double[]{10, 20, 50, 100, 200, 500, 1000, 2000}) {
        QuantileEstimator sq = new QuantileEstimator(1001);
        TDigest dist = new MergingDigest(compression);
        double[] data = new double[100000];
        for (int i = 0; i < 100000; i++) {
            double x = gen.nextDouble();
            dist.add(x);
            sq.add(x);
            data[i] = x;
        }
        dist.compress();
        Arrays.sort(data);

        List<Double> qz = sq.getQuantiles();
        for (double q : quantiles) {
            double x1 = dist.quantile(q);
            double x2 = qz.get((int) (q * 1000 + 0.5));
            double e1 = Dist.cdf(x1, data) - q;
            double e2 = Dist.cdf(x2, data) - q;
            out.printf("%s,%.0f,%.8f,%.10g,%.10g,%d,%d\n",
                    tag, compression, q, e1, e2, dist.smallByteSize(), sq.serializedSize());

        }
    }
}
 

private void testAccuracyInternal(double quantile, int size, double maxDelta, TDigest tdigest) {
  Random rand = new Random(12345);
  ScalableStatistics approx = 
      tdigest == null ? 
          new ScalableStatistics(0) : 
          new ScalableStatistics(0, tdigest);
  ScalableStatistics exact  = new ScalableStatistics(size + 3);
  for (int i = 0; i < size; i++) {
    double value = rand.nextDouble();
    exact.add(value);
    approx.add(value);
    assertEquals(exact.getQuantile(quantile), approx.getQuantile(quantile), maxDelta);
  }
}
 

@Override
public void serialize(TDigest td, JsonGenerator j, SerializerProvider sp) throws IOException, JsonProcessingException {
    try(ByteArrayOutputStream baos = new ByteArrayOutputStream(); ObjectOutputStream oos = new ObjectOutputStream(baos)){
        oos.writeObject(td);
        oos.close();
        byte[] bytes = baos.toByteArray();
        Base64 b = new Base64();
        String str = b.encodeAsString(bytes);
        j.writeStartObject();
        j.writeStringField("digest", str);
        j.writeEndObject();
    }
}
 

public static ObjectType getObjectType(Object value) {
  if (value instanceof String) {
    return ObjectType.String;
  } else if (value instanceof Long) {
    return ObjectType.Long;
  } else if (value instanceof Double) {
    return ObjectType.Double;
  } else if (value instanceof DoubleArrayList) {
    return ObjectType.DoubleArrayList;
  } else if (value instanceof AvgPair) {
    return ObjectType.AvgPair;
  } else if (value instanceof MinMaxRangePair) {
    return ObjectType.MinMaxRangePair;
  } else if (value instanceof HyperLogLog) {
    return ObjectType.HyperLogLog;
  } else if (value instanceof QuantileDigest) {
    return ObjectType.QuantileDigest;
  } else if (value instanceof Map) {
    return ObjectType.Map;
  } else if (value instanceof IntSet) {
    return ObjectType.IntSet;
  } else if (value instanceof TDigest) {
    return ObjectType.TDigest;
  } else if (value instanceof DistinctTable) {
    return ObjectType.DistinctTable;
  } else if (value instanceof Sketch) {
    return ObjectType.DataSketch;
  } else {
    throw new IllegalArgumentException("Unsupported type of value: " + value.getClass().getSimpleName());
  }
}
 

/**
 * Build Avro file containing serialized TDigest bytes.
 *
 * @param schema Schema of data (one fixed and one variable column)
 * @param _fixedExpected Serialized bytes of fixed length column are populated here
 * @param _varExpected Serialized bytes of variable length column are populated here
 * @throws IOException
 */
private void buildAvro(Schema schema, List<byte[]> _fixedExpected, List<byte[]> _varExpected)
    throws IOException {
  org.apache.avro.Schema avroSchema = AvroUtils.getAvroSchemaFromPinotSchema(schema);

  try (DataFileWriter<GenericData.Record> recordWriter = new DataFileWriter<>(new GenericDatumWriter<>(avroSchema))) {

    if (!new File(AVRO_DIR_NAME).mkdir()) {
      throw new RuntimeException("Unable to create test directory: " + AVRO_DIR_NAME);
    }

    recordWriter.create(avroSchema, new File(AVRO_DIR_NAME, AVRO_NAME));
    for (int i = 0; i < NUM_ROWS; i++) {
      GenericData.Record record = new GenericData.Record(avroSchema);

      TDigest tDigest = TDigest.createMergingDigest(PercentileTDigestAggregationFunction.DEFAULT_TDIGEST_COMPRESSION);
      tDigest.add(_random.nextDouble());

      ByteBuffer buffer = ByteBuffer.allocate(tDigest.byteSize());
      tDigest.asBytes(buffer);
      _fixedExpected.add(buffer.array());

      buffer.flip();
      record.put(FIXED_BYTES_UNSORTED_COLUMN, buffer);

      if (i % 2 == 0) {
        tDigest.add(_random.nextDouble());
      }

      buffer = ByteBuffer.allocate(tDigest.byteSize());
      tDigest.asBytes(buffer);
      _varExpected.add(buffer.array());

      buffer.flip();
      record.put(VARIABLE_BYTES_COLUMN, buffer);

      recordWriter.append(record);
    }
  }
}
 

@Test
public void testInnerSegmentAggregation() {
  // For inner segment case, percentile does not affect the intermediate result
  AggregationOperator aggregationOperator = getOperatorForQuery(getAggregationQuery(0));
  IntermediateResultsBlock resultsBlock = aggregationOperator.nextBlock();
  List<Object> aggregationResult = resultsBlock.getAggregationResult();
  Assert.assertNotNull(aggregationResult);
  Assert.assertEquals(aggregationResult.size(), 3);
  DoubleList doubleList = (DoubleList) aggregationResult.get(0);
  Collections.sort(doubleList);
  assertTDigest((TDigest) aggregationResult.get(1), doubleList);
  assertTDigest((TDigest) aggregationResult.get(2), doubleList);
}
 

private void assertTDigest(TDigest tDigest, DoubleList doubleList) {
  for (int percentile = 0; percentile <= 100; percentile++) {
    double expected;
    if (percentile == 100) {
      expected = doubleList.getDouble(doubleList.size() - 1);
    } else {
      expected = doubleList.getDouble(doubleList.size() * percentile / 100);
    }
    Assert.assertEquals(tDigest.quantile(percentile / 100.0), expected, DELTA, ERROR_MESSAGE);
  }
}
 

@Override
Object getRandomRawValue(Random random) {
  TDigest tDigest = TDigest.createMergingDigest(COMPRESSION);
  tDigest.add(random.nextLong());
  tDigest.add(random.nextLong());
  return ObjectSerDeUtils.TDIGEST_SER_DE.serialize(tDigest);
}