<commit_before>package com.firstblick.elasticsearch.action.export;

import org.elasticsearch.common.xcontent.ToXContent;
import org.elasticsearch.common.xcontent.XContentBuilder;
import org.elasticsearch.rest.action.support.RestXContentBuilder;
import org.elasticsearch.search.SearchHitField;
import org.elasticsearch.search.internal.InternalSearchHit;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class ExportFields implements ToXContent {

    private final List<String> fields;
    private InternalSearchHit hit;
    private final List<FieldExtractor> fieldExtractors;

    abstract class FieldExtractor implements ToXContent {
    }

    class HitFieldExtractor extends FieldExtractor {

        private final String fieldName;

        public HitFieldExtractor(String fieldName) {
            this.fieldName = fieldName;
        }

        @Override
        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
            SearchHitField field = hit.getFields().get(fieldName);
            if (field != null && !field.values().isEmpty()) {
                if (field.values().size() == 1) {
                    builder.field(field.name(), field.values().get(0));
                } else {
                    builder.field(field.name());
                    builder.startArray();
                    for (Object value : field.values()) {
                        builder.value(value);
                    }
                    builder.endArray();
                }
            }
            return builder;
        }
    }

    public void hit(InternalSearchHit hit) {
        this.hit = hit;
    }

    public ExportFields(List<String> fields) {
        this.fields = fields;
        this.fieldExtractors = getFieldExtractors();
    }

    private List<FieldExtractor> getFieldExtractors() {
        List<FieldExtractor> extractors = new ArrayList<FieldExtractor>(fields
                .size());
        for (String fn : fields) {
            FieldExtractor fc = null;
            if (fn.startsWith("_")) {
                if (fn.equals("_source")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            RestXContentBuilder.restDocumentSource(hit.sourceRef(), builder, params);
                            return builder;
                        }
                    };
                } else if (fn.equals("_id")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field("_id", hit.getId());
                        }
                    };
                } else if (fn.equals("_type")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field("_type", hit.getType());
                        }
                    };<commit_msg>reformat json of source upon export<commit_after>package com.firstblick.elasticsearch.action.export;

import org.elasticsearch.common.bytes.BytesReference;
import org.elasticsearch.common.xcontent.*;
import org.elasticsearch.search.SearchHitField;
import org.elasticsearch.search.internal.InternalSearchHit;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class ExportFields implements ToXContent {

    private final List<String> fields;
    private InternalSearchHit hit;
    private final List<FieldExtractor> fieldExtractors;

    abstract class FieldExtractor implements ToXContent {
    }

    class SourceFieldExtractor extends FieldExtractor {

        @Override
        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
            BytesReference source = hit.sourceRef();
            XContentType contentType = XContentFactory.xContentType(source);
            XContentParser parser = XContentFactory.xContent(contentType)
                    .createParser(source);
            try {
                parser.nextToken();
                builder.field("_source");
                builder.copyCurrentStructure(parser);
            } finally {
                parser.close();
            }
            return builder;
        }
    }

    class HitFieldExtractor extends FieldExtractor {

        private final String fieldName;

        public HitFieldExtractor(String fieldName) {
            this.fieldName = fieldName;
        }

        @Override
        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
            SearchHitField field = hit.getFields().get(fieldName);
            if (field != null && !field.values().isEmpty()) {
                if (field.values().size() == 1) {
                    builder.field(field.name(), field.values().get(0));
                } else {
                    builder.field(field.name());
                    builder.startArray();
                    for (Object value : field.values()) {
                        builder.value(value);
                    }
                    builder.endArray();
                }
            }
            return builder;
        }
    }

    public void hit(InternalSearchHit hit) {
        this.hit = hit;
    }

    public ExportFields(List<String> fields) {
        this.fields = fields;
        this.fieldExtractors = getFieldExtractors();
    }

    private List<FieldExtractor> getFieldExtractors() {
        List<FieldExtractor> extractors = new ArrayList<FieldExtractor>(fields
                .size());
        for (String fn : fields) {
            FieldExtractor fc = null;
            if (fn.startsWith("_")) {
                if (fn.equals("_source")) {
                    fc = new SourceFieldExtractor();
                } else if (fn.equals("_id")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field("_id", hit.getId());
                        }
                    };
                } else if (fn.equals("_type")) {
                    fc = new FieldExtractor() {
                        @Override
                        public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
                            return builder.field("_type", hit.getType());
                        }
                    };
{'fraction_non_alphanumeric': 0.08828671328671328, 'fraction_numerical': 0.0005827505827505828, 'mean_word_length': 1.7188118811881188, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, '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': '3874844', 'n_tokens_mistral': 1640, 'n_tokens_neox': 1572, 'n_words': 358}