Examples with BigArrays org.opensearch.common.util.BigArrays used on opensource projects

Example 1 with BigArrays

use of org.opensearch.common.util.BigArrays in project OpenSearch by opensearch-project.

the class MatrixStatsAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    if (valuesSources == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final BigArrays bigArrays = context.bigArrays();
    final NumericDoubleValues[] values = new NumericDoubleValues[valuesSources.fieldNames().length];
    for (int i = 0; i < values.length; ++i) {
        values[i] = valuesSources.getField(i, ctx);
    }
    return new LeafBucketCollectorBase(sub, values) {

        final String[] fieldNames = valuesSources.fieldNames();

        final double[] fieldVals = new double[fieldNames.length];

        @Override
        public void collect(int doc, long bucket) throws IOException {
            // get fields
            if (includeDocument(doc)) {
                stats = bigArrays.grow(stats, bucket + 1);
                RunningStats stat = stats.get(bucket);
                // add document fields to correlation stats
                if (stat == null) {
                    stat = new RunningStats(fieldNames, fieldVals);
                    stats.set(bucket, stat);
                } else {
                    stat.add(fieldNames, fieldVals);
                }
            }
        }

        /**
         * return a map of field names and data
         */
        private boolean includeDocument(int doc) throws IOException {
            // loop over fields
            for (int i = 0; i < fieldVals.length; ++i) {
                final NumericDoubleValues doubleValues = values[i];
                if (doubleValues.advanceExact(doc)) {
                    final double value = doubleValues.doubleValue();
                    if (value == Double.NEGATIVE_INFINITY) {
                        // TODO: Fix matrix stats to treat neg inf as any other value
                        return false;
                    }
                    fieldVals[i] = value;
                } else {
                    return false;
                }
            }
            return true;
        }
    };
}

Example 2 with BigArrays

use of org.opensearch.common.util.BigArrays in project OpenSearch by opensearch-project.

the class TermsValuesSourceBuilder method register.

static void register(ValuesSourceRegistry.Builder builder) {
    builder.register(REGISTRY_KEY, org.opensearch.common.collect.List.of(CoreValuesSourceType.DATE, CoreValuesSourceType.NUMERIC, CoreValuesSourceType.BOOLEAN), (valuesSourceConfig, name, hasScript, format, missingBucket, missingOrder, order) -> {
        final DocValueFormat docValueFormat;
        if (format == null && valuesSourceConfig.valueSourceType() == CoreValuesSourceType.DATE) {
            // defaults to the raw format on date fields (preserve timestamp as longs).
            docValueFormat = DocValueFormat.RAW;
        } else {
            docValueFormat = valuesSourceConfig.format();
        }
        return new CompositeValuesSourceConfig(name, valuesSourceConfig.fieldType(), valuesSourceConfig.getValuesSource(), docValueFormat, order, missingBucket, missingOrder, hasScript, (BigArrays bigArrays, IndexReader reader, int size, LongConsumer addRequestCircuitBreakerBytes, CompositeValuesSourceConfig compositeValuesSourceConfig) -> {
            final ValuesSource.Numeric vs = (ValuesSource.Numeric) compositeValuesSourceConfig.valuesSource();
            if (vs.isFloatingPoint()) {
                return new DoubleValuesSource(bigArrays, compositeValuesSourceConfig.fieldType(), vs::doubleValues, compositeValuesSourceConfig.format(), compositeValuesSourceConfig.missingBucket(), compositeValuesSourceConfig.missingOrder(), size, compositeValuesSourceConfig.reverseMul());
            } else {
                final LongUnaryOperator rounding;
                rounding = LongUnaryOperator.identity();
                return new LongValuesSource(bigArrays, compositeValuesSourceConfig.fieldType(), vs::longValues, rounding, compositeValuesSourceConfig.format(), compositeValuesSourceConfig.missingBucket(), compositeValuesSourceConfig.missingOrder(), size, compositeValuesSourceConfig.reverseMul());
            }
        });
    }, false);
    builder.register(REGISTRY_KEY, org.opensearch.common.collect.List.of(CoreValuesSourceType.BYTES, CoreValuesSourceType.IP), (valuesSourceConfig, name, hasScript, format, missingBucket, missingOrder, order) -> new CompositeValuesSourceConfig(name, valuesSourceConfig.fieldType(), valuesSourceConfig.getValuesSource(), valuesSourceConfig.format(), order, missingBucket, missingOrder, hasScript, (BigArrays bigArrays, IndexReader reader, int size, LongConsumer addRequestCircuitBreakerBytes, CompositeValuesSourceConfig compositeValuesSourceConfig) -> {
        if (valuesSourceConfig.hasGlobalOrdinals() && reader instanceof DirectoryReader) {
            ValuesSource.Bytes.WithOrdinals vs = (ValuesSource.Bytes.WithOrdinals) compositeValuesSourceConfig.valuesSource();
            return new GlobalOrdinalValuesSource(bigArrays, compositeValuesSourceConfig.fieldType(), vs::globalOrdinalsValues, compositeValuesSourceConfig.format(), compositeValuesSourceConfig.missingBucket(), compositeValuesSourceConfig.missingOrder(), size, compositeValuesSourceConfig.reverseMul());
        } else {
            ValuesSource.Bytes vs = (ValuesSource.Bytes) compositeValuesSourceConfig.valuesSource();
            return new BinaryValuesSource(bigArrays, addRequestCircuitBreakerBytes, compositeValuesSourceConfig.fieldType(), vs::bytesValues, compositeValuesSourceConfig.format(), compositeValuesSourceConfig.missingBucket(), compositeValuesSourceConfig.missingOrder(), size, compositeValuesSourceConfig.reverseMul());
        }
    }), false);
}

Example 3 with BigArrays

use of org.opensearch.common.util.BigArrays in project OpenSearch by opensearch-project.

the class GeoTileGridValuesSourceBuilder method register.

static void register(ValuesSourceRegistry.Builder builder) {
    builder.register(REGISTRY_KEY, CoreValuesSourceType.GEOPOINT, (valuesSourceConfig, precision, boundingBox, name, hasScript, format, missingBucket, missingOrder, order) -> {
        ValuesSource.GeoPoint geoPoint = (ValuesSource.GeoPoint) valuesSourceConfig.getValuesSource();
        // is specified in the builder.
        final MappedFieldType fieldType = valuesSourceConfig.fieldType();
        CellIdSource cellIdSource = new CellIdSource(geoPoint, precision, boundingBox, GeoTileUtils::longEncode);
        return new CompositeValuesSourceConfig(name, fieldType, cellIdSource, DocValueFormat.GEOTILE, order, missingBucket, missingOrder, hasScript, (BigArrays bigArrays, IndexReader reader, int size, LongConsumer addRequestCircuitBreakerBytes, CompositeValuesSourceConfig compositeValuesSourceConfig) -> {
            final CellIdSource cis = (CellIdSource) compositeValuesSourceConfig.valuesSource();
            return new GeoTileValuesSource(bigArrays, compositeValuesSourceConfig.fieldType(), cis::longValues, LongUnaryOperator.identity(), compositeValuesSourceConfig.format(), compositeValuesSourceConfig.missingBucket(), compositeValuesSourceConfig.missingOrder(), size, compositeValuesSourceConfig.reverseMul());
        });
    }, false);
}

Example 4 with BigArrays

use of org.opensearch.common.util.BigArrays in project OpenSearch by opensearch-project.

the class HistogramValuesSourceBuilder method register.

public static void register(ValuesSourceRegistry.Builder builder) {
    builder.register(REGISTRY_KEY, org.opensearch.common.collect.List.of(CoreValuesSourceType.DATE, CoreValuesSourceType.NUMERIC), (valuesSourceConfig, interval, name, hasScript, format, missingBucket, missingOrder, order) -> {
        ValuesSource.Numeric numeric = (ValuesSource.Numeric) valuesSourceConfig.getValuesSource();
        final HistogramValuesSource vs = new HistogramValuesSource(numeric, interval);
        final MappedFieldType fieldType = valuesSourceConfig.fieldType();
        return new CompositeValuesSourceConfig(name, fieldType, vs, valuesSourceConfig.format(), order, missingBucket, missingOrder, hasScript, (BigArrays bigArrays, IndexReader reader, int size, LongConsumer addRequestCircuitBreakerBytes, CompositeValuesSourceConfig compositeValuesSourceConfig) -> {
            final ValuesSource.Numeric numericValuesSource = (ValuesSource.Numeric) compositeValuesSourceConfig.valuesSource();
            return new DoubleValuesSource(bigArrays, compositeValuesSourceConfig.fieldType(), numericValuesSource::doubleValues, compositeValuesSourceConfig.format(), compositeValuesSourceConfig.missingBucket(), compositeValuesSourceConfig.missingOrder(), size, compositeValuesSourceConfig.reverseMul());
        });
    }, false);
}

Example 5 with BigArrays

use of org.opensearch.common.util.BigArrays in project OpenSearch by opensearch-project.

the class SumAggregator method getLeafCollector.

@Override
public LeafBucketCollector getLeafCollector(LeafReaderContext ctx, final LeafBucketCollector sub) throws IOException {
    if (valuesSource == null) {
        return LeafBucketCollector.NO_OP_COLLECTOR;
    }
    final BigArrays bigArrays = context.bigArrays();
    final SortedNumericDoubleValues values = valuesSource.doubleValues(ctx);
    final CompensatedSum kahanSummation = new CompensatedSum(0, 0);
    return new LeafBucketCollectorBase(sub, values) {

        @Override
        public void collect(int doc, long bucket) throws IOException {
            sums = bigArrays.grow(sums, bucket + 1);
            compensations = bigArrays.grow(compensations, bucket + 1);
            if (values.advanceExact(doc)) {
                final int valuesCount = values.docValueCount();
                // Compute the sum of double values with Kahan summation algorithm which is more
                // accurate than naive summation.
                double sum = sums.get(bucket);
                double compensation = compensations.get(bucket);
                kahanSummation.reset(sum, compensation);
                for (int i = 0; i < valuesCount; i++) {
                    double value = values.nextValue();
                    kahanSummation.add(value);
                }
                compensations.set(bucket, kahanSummation.delta());
                sums.set(bucket, kahanSummation.value());
            }
        }
    };
}