<commit_before>package com.jivesoftware.os.miru.catwalk.deployable;

import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.PeekingIterator;
import com.google.common.collect.Sets;
import com.jivesoftware.os.amza.api.PartitionClient;
import com.jivesoftware.os.amza.api.PartitionClientProvider;
import com.jivesoftware.os.amza.api.filer.UIO;
import com.jivesoftware.os.amza.api.partition.Consistency;
import com.jivesoftware.os.amza.api.partition.Durability;
import com.jivesoftware.os.amza.api.partition.PartitionName;
import com.jivesoftware.os.amza.api.partition.PartitionProperties;
import com.jivesoftware.os.amza.api.stream.RowType;
import com.jivesoftware.os.amza.service.filer.HeapFiler;
import com.jivesoftware.os.miru.api.MiruStats;
import com.jivesoftware.os.miru.api.base.MiruTenantId;
import com.jivesoftware.os.miru.api.base.MiruTermId;
import com.jivesoftware.os.miru.plugin.solution.MiruTimeRange;
import com.jivesoftware.os.miru.stream.plugins.catwalk.CatwalkModel;
import com.jivesoftware.os.miru.stream.plugins.catwalk.CatwalkQuery;
import com.jivesoftware.os.miru.stream.plugins.catwalk.FeatureScore;
import com.jivesoftware.os.mlogger.core.MetricLogger;
import com.jivesoftware.os.mlogger.core.MetricLoggerFactory;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.TreeSet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 *
 */
public class CatwalkModelService {

    private static final MetricLogger LOG = MetricLoggerFactory.getLogger();

    private static final PartitionProperties MODEL_PROPERTIES = new PartitionProperties(Durability.fsync_async,
        TimeUnit.DAYS.toMillis(30), TimeUnit.DAYS.toMillis(10), TimeUnit.DAYS.toMillis(30), TimeUnit.DAYS.toMillis(10), 0, 0, 0, 0,
        false, Consistency.leader_quorum, true, true, false, RowType.snappy_primary, "berkeleydb", null, -1, -1);

    private final CatwalkModelQueue modelQueue;
    private final ExecutorService readRepairers;
    private final PartitionClientProvider clientProvider;
    private final MiruStats stats;

    private final long additionalSolverAfterNMillis = 1_000; //TODO expose to conf?
    private final long abandonLeaderSolutionAfterNMillis = 5_000; //TODO expose to conf?
    private final long abandonSolutionAfterNMillis = 30_000; //TODO expose to conf?

    public CatwalkModelService(CatwalkModelQueue modelQueue,
        ExecutorService readRepairers,
        PartitionClientProvider clientProvider,
        MiruStats stats) {
        this.modelQueue = modelQueue;
        this.readRepairers = readRepairers;
        this.clientProvider = clientProvider;
        this.stats = stats;
    }

    public CatwalkModel getModel(MiruTenantId tenantId, String catwalkId, String modelId, CatwalkQuery catwalkQuery) throws Exception {
        long start = System.currentTimeMillis();

        PartitionClient client = modelClient(tenantId);

        List<FeatureRange> deletableRanges = Lists.newArrayList();
        TreeSet<Integer> partitionIds = Sets.newTreeSet();

        Map<FieldIdsKey, MergedScores> fieldIdsToFeatureScores = new HashMap<>();
        FeatureRange[] currentRange = { null };
        String[][] featureFields = catwalkQuery.featureFields;
        client.scan(Consistency.leader_quorum,
            prefixedKeyRangeStream -> {
                for (int i = 0; i < featureFields.length; i++) {
                    String[] fields = featureFields[i];
                    byte[] fromKey = modelPartitionKey(catwalkId, modelId, fields, 0, Integer.MAX_VALUE);
                    byte[] toKey = modelPartitionKey(catwalkId, modelId, fields, Integer.MAX_VALUE, 0);
                    if (!prefixedKeyRangeStream.stream(null, fromKey, null, toKey)) {
                        return false;
                    }
                }
                return true;
            },
            (prefix, key, value, timestamp, version) -> {

                FeatureRange range = getFeatureRange(key);
                if (currentRange[0] == null || !currentRange[0].intersects(range)) {
                    currentRange[0] = range;

                    for (int partitionId = range.fromPartitionId; partitionId <= range.toPartitionId; partitionId++) {
                        partitionIds.add(partitionId);
                    }

                    ModelFeatureScores scores = valueFromBytes(value);
                    FieldIdsKey fieldIdsKey = new FieldIdsKey(range.fieldIds);
                    fieldIdsToFeatureScores.compute(fieldIdsKey, (t, currentMerged) -> {
                        if (currentMerged == null) {
                            return new MergedScores(range, scores);
                        }

                        PeekingIterator<FeatureScore> a = Iterators.peekingIterator(currentMerged.latestScores.featureScores.iterator());
                        PeekingIterator<FeatureScore> b = Iterators.peekingIterator(scores.featureScores.iterator());

                        List<FeatureScore> merged = new ArrayList<>(currentMerged.latestScores.featureScores.size() + scores.featureScores.size());
                        while (a.hasNext() || b.hasNext()) {

                            if (a.hasNext() && b.hasNext()) {
                                int c = FEATURE_SCORE_COMPARATOR.compare(a.peek(), b.peek());
                                if (c == 0) {
                                    merged.add(merge(a.next(), b.next()));
                                } else if (c < 0) {
                                    merged.add(a.next());
                                } else {
                                    merged.add(b.next());
                                }
                            } else if (a.hasNext()) {
                                merged.add(a.next());

                            } else if (b.hasNext()) {
                                merged.add(b.next());
                            }
                        }

                        currentMerged.contiguousClosedPartitions &= scores.partitionIsClosed;
                        currentMerged.contiguousClosedPartitions &= (currentMerged.latestRange.toPartitionId + 1 == range.fromPartitionId);
                        if (currentMerged.contiguousClosedPartitions) {
                            if (currentMerged.ranges == null) {
                                currentMerged.ranges = new ArrayList<>();
                            }
                            currentMerged.ranges.add(currentMerged.latestRange);
                            currentMerged.scores = currentMerged.latestScores;
                            if (currentMerged.timeRange == null) {
                                currentMerged.timeRange = currentMerged.latestScores.timeRange;
                            } else {
                                currentMerged.timeRange = new MiruTimeRange(
                                    Math.min(scores.timeRange.smallestTimestamp, currentMerged.timeRange.smallestTimestamp),
                                    Math.max(scores.timeRange.largestTimestamp, currentMerged.timeRange.largestTimestamp));
                            }
                        }

                        currentMerged.numberOfMerges++;
                        currentMerged.latestRange = range;
                        currentMerged.latestScores = new ModelFeatureScores(scores.partitionIsClosed, merged, scores.timeRange);
                        return currentMerged;
                    });

                } else {
                    deletableRanges.add(range);
                }
                return true;
            },
            additionalSolverAfterNMillis,
            abandonLeaderSolutionAfterNMillis,
            abandonSolutionAfterNMillis,
            Optional.<List<String>>empty());

        @SuppressWarnings("unchecked")
        List<FeatureScore>[] featureScores = new List[featureFields.length];
        for (int i = 0; i < featureFields.length; i++) {
            MergedScores mergedScores = fieldIdsToFeatureScores.get(new FieldIdsKey(featureFields[i]));
            featureScores[i] = mergedScores.latestScores.featureScores;
            LOG.info("Gathered {} scores for tenantId:{} catwalkId:{} modelId:{} feature:{} from {} models",
                featureScores[i].size(), tenantId, catwalkId, modelId, i, 1 + mergedScores.numberOfMerges);
        }

        for (Map.Entry<FieldIdsKey, MergedScores> entry : fieldIdsToFeatureScores.entrySet()) {
            List<FeatureRange> ranges = entry.getValue().ranges;
            if (ranges != null && ranges.size() > 1) {
                readRepairers.submit(new ReadRepair(tenantId, catwalkId, modelId, entry.getKey(), entry.getValue()));
            }
        }

        if (!deletableRanges.isEmpty()) {
            readRepairers.submit(() -> {
                try {
                    removeModel(tenantId, catwalkId, modelId, deletableRanges);
                } catch (Exception x) {
                    LOG.error("Failure while trying to delete.");
                }
            });<commit_msg>fix NPE when gathering catwalk model with missing feature<commit_after>package com.jivesoftware.os.miru.catwalk.deployable;

import com.google.common.collect.Iterators;
import com.google.common.collect.Lists;
import com.google.common.collect.PeekingIterator;
import com.google.common.collect.Sets;
import com.jivesoftware.os.amza.api.PartitionClient;
import com.jivesoftware.os.amza.api.PartitionClientProvider;
import com.jivesoftware.os.amza.api.filer.UIO;
import com.jivesoftware.os.amza.api.partition.Consistency;
import com.jivesoftware.os.amza.api.partition.Durability;
import com.jivesoftware.os.amza.api.partition.PartitionName;
import com.jivesoftware.os.amza.api.partition.PartitionProperties;
import com.jivesoftware.os.amza.api.stream.RowType;
import com.jivesoftware.os.amza.service.filer.HeapFiler;
import com.jivesoftware.os.miru.api.MiruStats;
import com.jivesoftware.os.miru.api.base.MiruTenantId;
import com.jivesoftware.os.miru.api.base.MiruTermId;
import com.jivesoftware.os.miru.plugin.solution.MiruTimeRange;
import com.jivesoftware.os.miru.stream.plugins.catwalk.CatwalkModel;
import com.jivesoftware.os.miru.stream.plugins.catwalk.CatwalkQuery;
import com.jivesoftware.os.miru.stream.plugins.catwalk.FeatureScore;
import com.jivesoftware.os.mlogger.core.MetricLogger;
import com.jivesoftware.os.mlogger.core.MetricLoggerFactory;
import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.TreeSet;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.TimeUnit;

/**
 *
 */
public class CatwalkModelService {

    private static final MetricLogger LOG = MetricLoggerFactory.getLogger();

    private static final PartitionProperties MODEL_PROPERTIES = new PartitionProperties(Durability.fsync_async,
        TimeUnit.DAYS.toMillis(30), TimeUnit.DAYS.toMillis(10), TimeUnit.DAYS.toMillis(30), TimeUnit.DAYS.toMillis(10), 0, 0, 0, 0,
        false, Consistency.leader_quorum, true, true, false, RowType.snappy_primary, "berkeleydb", null, -1, -1);

    private final CatwalkModelQueue modelQueue;
    private final ExecutorService readRepairers;
    private final PartitionClientProvider clientProvider;
    private final MiruStats stats;

    private final long additionalSolverAfterNMillis = 1_000; //TODO expose to conf?
    private final long abandonLeaderSolutionAfterNMillis = 5_000; //TODO expose to conf?
    private final long abandonSolutionAfterNMillis = 30_000; //TODO expose to conf?

    public CatwalkModelService(CatwalkModelQueue modelQueue,
        ExecutorService readRepairers,
        PartitionClientProvider clientProvider,
        MiruStats stats) {
        this.modelQueue = modelQueue;
        this.readRepairers = readRepairers;
        this.clientProvider = clientProvider;
        this.stats = stats;
    }

    public CatwalkModel getModel(MiruTenantId tenantId, String catwalkId, String modelId, CatwalkQuery catwalkQuery) throws Exception {
        long start = System.currentTimeMillis();

        PartitionClient client = modelClient(tenantId);

        List<FeatureRange> deletableRanges = Lists.newArrayList();
        TreeSet<Integer> partitionIds = Sets.newTreeSet();

        Map<FieldIdsKey, MergedScores> fieldIdsToFeatureScores = new HashMap<>();
        FeatureRange[] currentRange = { null };
        String[][] featureFields = catwalkQuery.featureFields;
        client.scan(Consistency.leader_quorum,
            prefixedKeyRangeStream -> {
                for (int i = 0; i < featureFields.length; i++) {
                    String[] fields = featureFields[i];
                    byte[] fromKey = modelPartitionKey(catwalkId, modelId, fields, 0, Integer.MAX_VALUE);
                    byte[] toKey = modelPartitionKey(catwalkId, modelId, fields, Integer.MAX_VALUE, 0);
                    if (!prefixedKeyRangeStream.stream(null, fromKey, null, toKey)) {
                        return false;
                    }
                }
                return true;
            },
            (prefix, key, value, timestamp, version) -> {

                FeatureRange range = getFeatureRange(key);
                if (currentRange[0] == null || !currentRange[0].intersects(range)) {
                    currentRange[0] = range;

                    for (int partitionId = range.fromPartitionId; partitionId <= range.toPartitionId; partitionId++) {
                        partitionIds.add(partitionId);
                    }

                    ModelFeatureScores scores = valueFromBytes(value);
                    FieldIdsKey fieldIdsKey = new FieldIdsKey(range.fieldIds);
                    fieldIdsToFeatureScores.compute(fieldIdsKey, (t, currentMerged) -> {
                        if (currentMerged == null) {
                            return new MergedScores(range, scores);
                        }

                        PeekingIterator<FeatureScore> a = Iterators.peekingIterator(currentMerged.latestScores.featureScores.iterator());
                        PeekingIterator<FeatureScore> b = Iterators.peekingIterator(scores.featureScores.iterator());

                        List<FeatureScore> merged = new ArrayList<>(currentMerged.latestScores.featureScores.size() + scores.featureScores.size());
                        while (a.hasNext() || b.hasNext()) {

                            if (a.hasNext() && b.hasNext()) {
                                int c = FEATURE_SCORE_COMPARATOR.compare(a.peek(), b.peek());
                                if (c == 0) {
                                    merged.add(merge(a.next(), b.next()));
                                } else if (c < 0) {
                                    merged.add(a.next());
                                } else {
                                    merged.add(b.next());
                                }
                            } else if (a.hasNext()) {
                                merged.add(a.next());

                            } else if (b.hasNext()) {
                                merged.add(b.next());
                            }
                        }

                        currentMerged.contiguousClosedPartitions &= scores.partitionIsClosed;
                        currentMerged.contiguousClosedPartitions &= (currentMerged.latestRange.toPartitionId + 1 == range.fromPartitionId);
                        if (currentMerged.contiguousClosedPartitions) {
                            if (currentMerged.ranges == null) {
                                currentMerged.ranges = new ArrayList<>();
                            }
                            currentMerged.ranges.add(currentMerged.latestRange);
                            currentMerged.scores = currentMerged.latestScores;
                            if (currentMerged.timeRange == null) {
                                currentMerged.timeRange = currentMerged.latestScores.timeRange;
                            } else {
                                currentMerged.timeRange = new MiruTimeRange(
                                    Math.min(scores.timeRange.smallestTimestamp, currentMerged.timeRange.smallestTimestamp),
                                    Math.max(scores.timeRange.largestTimestamp, currentMerged.timeRange.largestTimestamp));
                            }
                        }

                        currentMerged.numberOfMerges++;
                        currentMerged.latestRange = range;
                        currentMerged.latestScores = new ModelFeatureScores(scores.partitionIsClosed, merged, scores.timeRange);
                        return currentMerged;
                    });

                } else {
                    deletableRanges.add(range);
                }
                return true;
            },
            additionalSolverAfterNMillis,
            abandonLeaderSolutionAfterNMillis,
            abandonSolutionAfterNMillis,
            Optional.<List<String>>empty());

        @SuppressWarnings("unchecked")
        List<FeatureScore>[] featureScores = new List[featureFields.length];
        for (int i = 0; i < featureFields.length; i++) {
            MergedScores mergedScores = fieldIdsToFeatureScores.get(new FieldIdsKey(featureFields[i]));
            if (mergedScores != null) {
                featureScores[i] = mergedScores.latestScores.featureScores;
                LOG.info("Gathered {} scores for tenantId:{} catwalkId:{} modelId:{} feature:{} from {} models",
                    featureScores[i].size(), tenantId, catwalkId, modelId, i, 1 + mergedScores.numberOfMerges);
            } else {
                featureScores[i] = Collections.emptyList();
                LOG.info("Gathered no scores for tenantId:{} catwalkId:{} modelId:{} feature:{}",
                    featureScores[i].size(), tenantId, catwalkId, modelId, i);
            }
        }

        for (Map.Entry<FieldIdsKey, MergedScores> entry : fieldIdsToFeatureScores.entrySet()) {
            List<FeatureRange> ranges = entry.getValue().ranges;
            if (ranges != null && ranges.size() > 1) {
                readRepairers.submit(new ReadRepair(tenantId, catwalkId, modelId, entry.getKey(), entry.getValue()));
            }
        }

        if (!deletableRanges.isEmpty()) {
            readRepairers.submit(() -> {
                try {
                    removeModel(tenantId, catwalkId, modelId, deletableRanges);
                } catch (Exception x) {
                    LOG.error("Failure while trying to delete.");
                }
            });
{'fraction_non_alphanumeric': 0.09732704402515724, 'fraction_numerical': 0.004088050314465409, 'mean_word_length': 2.047596230634084, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 18, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3442509', 'n_tokens_mistral': 5268, 'n_tokens_neox': 5043, 'n_words': 893}