<commit_before>package com.yahoo.search.dispatch.rpc;

import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol;
import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol.StringProperty;
import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol.TensorProperty;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import com.yahoo.data.access.simple.Value;
import com.yahoo.data.access.slime.SlimeAdapter;
import com.yahoo.document.GlobalId;
import com.yahoo.fs4.GetDocSumsPacket;
import com.yahoo.io.GrowableByteBuffer;
import com.yahoo.prelude.fastsearch.DocumentDatabase;
import com.yahoo.prelude.fastsearch.FastHit;
import com.yahoo.prelude.fastsearch.GroupingListHit;
import com.yahoo.prelude.fastsearch.VespaBackEndSearcher;
import com.yahoo.search.Query;
import com.yahoo.search.Result;
import com.yahoo.search.grouping.vespa.GroupingExecutor;
import com.yahoo.search.query.Model;
import com.yahoo.search.query.QueryTree;
import com.yahoo.search.query.Ranking;
import com.yahoo.search.query.Sorting;
import com.yahoo.search.query.Sorting.Order;
import com.yahoo.search.result.Coverage;
import com.yahoo.search.result.Relevance;
import com.yahoo.searchlib.aggregation.Grouping;
import com.yahoo.slime.BinaryFormat;
import com.yahoo.vespa.objects.BufferSerializer;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

public class ProtobufSerialization {

    private static final int INITIAL_SERIALIZATION_BUFFER_SIZE = 10 * 1024;

    public static byte[] serializeSearchRequest(Query query, String serverId) {
        return convertFromQuery(query, serverId).toByteArray();
    }

    private static SearchProtocol.SearchRequest convertFromQuery(Query query, String serverId) {
        var builder = SearchProtocol.SearchRequest.newBuilder().setHits(query.getHits()).setOffset(query.getOffset())
                .setTimeout((int) query.getTimeLeft());

        var documentDb = query.getModel().getDocumentDb();
        if (documentDb != null) {
            builder.setDocumentType(documentDb);
        }
        builder.setQueryTreeBlob(serializeQueryTree(query.getModel().getQueryTree()));

        if (query.getGroupingSessionCache() || query.getRanking().getQueryCache()) {
            // TODO verify that the session key is included whenever rank properties would have been
            builder.setSessionKey(query.getSessionId(serverId).toString());
        }
        if (query.properties().getBoolean(Model.ESTIMATE)) {
            builder.setHits(0);
        }
        if (GroupingExecutor.hasGroupingList(query)) {
            List<Grouping> groupingList = GroupingExecutor.getGroupingList(query);
            BufferSerializer gbuf = new BufferSerializer(new GrowableByteBuffer());
            gbuf.putInt(null, groupingList.size());
            for (Grouping g : groupingList) {
                g.serialize(gbuf);
            }
            gbuf.getBuf().flip();
            builder.setGroupingBlob(ByteString.copyFrom(gbuf.getBuf().getByteBuffer()));
        }
        if (query.getGroupingSessionCache()) {
            builder.setCacheGrouping(true);
        }

        mergeToSearchRequestFromRanking(query.getRanking(), builder);

        return builder.build();
    }

    private static void mergeToSearchRequestFromRanking(Ranking ranking, SearchProtocol.SearchRequest.Builder builder) {
        builder.setRankProfile(ranking.getProfile());

        if (ranking.getQueryCache()) {
            builder.setCacheQuery(true);
        }
        if (ranking.getSorting() != null) {
            mergeToSearchRequestFromSorting(ranking.getSorting(), builder);
        }
        if (ranking.getLocation() != null) {
            builder.setGeoLocation(ranking.getLocation().toString());
        }

        var featureMap = ranking.getFeatures().asMap();
        MapConverter.convertMapPrimitives(featureMap, builder::addFeatureOverrides);
        MapConverter.convertMapTensors(featureMap, builder::addTensorFeatureOverrides);
        mergeRankProperties(ranking, builder::addRankProperties, builder::addTensorRankProperties);
    }

    private static void mergeToSearchRequestFromSorting(Sorting sorting, SearchProtocol.SearchRequest.Builder builder) {
        for (var field : sorting.fieldOrders()) {
            var sortField = SearchProtocol.SortField.newBuilder()
                    .setField(field.getSorter().getName())
                    .setAscending(field.getSortOrder() == Order.ASCENDING).build();
            builder.addSorting(sortField);<commit_msg>Also need the sort function.<commit_after>package com.yahoo.search.dispatch.rpc;

import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol;
import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol.StringProperty;
import ai.vespa.searchlib.searchprotocol.protobuf.SearchProtocol.TensorProperty;
import com.google.protobuf.ByteString;
import com.google.protobuf.InvalidProtocolBufferException;
import com.yahoo.data.access.simple.Value;
import com.yahoo.data.access.slime.SlimeAdapter;
import com.yahoo.document.GlobalId;
import com.yahoo.fs4.GetDocSumsPacket;
import com.yahoo.io.GrowableByteBuffer;
import com.yahoo.prelude.fastsearch.DocumentDatabase;
import com.yahoo.prelude.fastsearch.FastHit;
import com.yahoo.prelude.fastsearch.GroupingListHit;
import com.yahoo.prelude.fastsearch.VespaBackEndSearcher;
import com.yahoo.search.Query;
import com.yahoo.search.Result;
import com.yahoo.search.grouping.vespa.GroupingExecutor;
import com.yahoo.search.query.Model;
import com.yahoo.search.query.QueryTree;
import com.yahoo.search.query.Ranking;
import com.yahoo.search.query.Sorting;
import com.yahoo.search.query.Sorting.Order;
import com.yahoo.search.result.Coverage;
import com.yahoo.search.result.Relevance;
import com.yahoo.searchlib.aggregation.Grouping;
import com.yahoo.slime.BinaryFormat;
import com.yahoo.vespa.objects.BufferSerializer;

import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.function.Consumer;

public class ProtobufSerialization {

    private static final int INITIAL_SERIALIZATION_BUFFER_SIZE = 10 * 1024;

    public static byte[] serializeSearchRequest(Query query, String serverId) {
        return convertFromQuery(query, serverId).toByteArray();
    }

    private static SearchProtocol.SearchRequest convertFromQuery(Query query, String serverId) {
        var builder = SearchProtocol.SearchRequest.newBuilder().setHits(query.getHits()).setOffset(query.getOffset())
                .setTimeout((int) query.getTimeLeft());

        var documentDb = query.getModel().getDocumentDb();
        if (documentDb != null) {
            builder.setDocumentType(documentDb);
        }
        builder.setQueryTreeBlob(serializeQueryTree(query.getModel().getQueryTree()));

        if (query.getGroupingSessionCache() || query.getRanking().getQueryCache()) {
            // TODO verify that the session key is included whenever rank properties would have been
            builder.setSessionKey(query.getSessionId(serverId).toString());
        }
        if (query.properties().getBoolean(Model.ESTIMATE)) {
            builder.setHits(0);
        }
        if (GroupingExecutor.hasGroupingList(query)) {
            List<Grouping> groupingList = GroupingExecutor.getGroupingList(query);
            BufferSerializer gbuf = new BufferSerializer(new GrowableByteBuffer());
            gbuf.putInt(null, groupingList.size());
            for (Grouping g : groupingList) {
                g.serialize(gbuf);
            }
            gbuf.getBuf().flip();
            builder.setGroupingBlob(ByteString.copyFrom(gbuf.getBuf().getByteBuffer()));
        }
        if (query.getGroupingSessionCache()) {
            builder.setCacheGrouping(true);
        }

        mergeToSearchRequestFromRanking(query.getRanking(), builder);

        return builder.build();
    }

    private static void mergeToSearchRequestFromRanking(Ranking ranking, SearchProtocol.SearchRequest.Builder builder) {
        builder.setRankProfile(ranking.getProfile());

        if (ranking.getQueryCache()) {
            builder.setCacheQuery(true);
        }
        if (ranking.getSorting() != null) {
            mergeToSearchRequestFromSorting(ranking.getSorting(), builder);
        }
        if (ranking.getLocation() != null) {
            builder.setGeoLocation(ranking.getLocation().toString());
        }

        var featureMap = ranking.getFeatures().asMap();
        MapConverter.convertMapPrimitives(featureMap, builder::addFeatureOverrides);
        MapConverter.convertMapTensors(featureMap, builder::addTensorFeatureOverrides);
        mergeRankProperties(ranking, builder::addRankProperties, builder::addTensorRankProperties);
    }

    private static void mergeToSearchRequestFromSorting(Sorting sorting, SearchProtocol.SearchRequest.Builder builder) {
        for (var field : sorting.fieldOrders()) {
            var sortField = SearchProtocol.SortField.newBuilder()
                    .setField(field.getSorter().toString())
                    .setAscending(field.getSortOrder() == Order.ASCENDING).build();
            builder.addSorting(sortField);
{'fraction_non_alphanumeric': 0.10679506600302965, 'fraction_numerical': 0.0017312270071413113, 'mean_word_length': 3.6966463414634148, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1797642', 'n_tokens_mistral': 2616, 'n_tokens_neox': 2498, 'n_words': 401}