<commit_before>package io.lumify.palantir.dataImport;

import io.lumify.core.model.properties.LumifyProperties;
import io.lumify.palantir.dataImport.model.PtMediaAndValue;
import org.securegraph.Vertex;
import org.securegraph.VertexBuilder;
import org.securegraph.property.StreamingPropertyValue;

import java.io.InputStream;
import java.util.zip.InflaterInputStream;

import static com.google.common.base.Preconditions.checkNotNull;

public class PtMediaAndValueImporter extends PtImporterBase<PtMediaAndValue> {
    public static final int ZLIB_COMPRESSED_TYPE_START = 2000; // TODO need to verify these numbers with other installations of Palantir
    public static final int ZLIB_COMPRESSED_TYPE_END = 3000;

    public PtMediaAndValueImporter(DataImporter dataImporter) {
        super(dataImporter, PtMediaAndValue.class);
    }

    @Override
    protected void beforeProcessRows() {
        super.beforeProcessRows();
        getDataImporter().getGraph().flush();
    }

    @Override
    protected void processRow(PtMediaAndValue row) throws Exception {
        String propertyKey = getDataImporter().getIdPrefix() + row.getId();
        InputStream in = row.getContents();
        try {
            if (row.getType() >= ZLIB_COMPRESSED_TYPE_START && row.getType() < ZLIB_COMPRESSED_TYPE_END) {
                in = new InflaterInputStream(in);
            }
            StreamingPropertyValue propertyValue = new StreamingPropertyValue(in, byte[].class);
            propertyValue.store(true);
            propertyValue.searchIndex(false);

            VertexBuilder vertexBuilder = getDataImporter().getGraph().prepareVertex(getMediaId(row), getDataImporter().getVisibility());
            LumifyProperties.RAW.addPropertyValue(vertexBuilder, propertyKey, propertyValue, getDataImporter().getVisibility());
            LumifyProperties.TITLE.setProperty(vertexBuilder, row.getTitle(), getDataImporter().getVisibility());
            Vertex mediaVertex = vertexBuilder.save(getDataImporter().getAuthorizations());

            Vertex sourceVertex = getVertexCache().get(getObjectId(row.getLinkObjectId()));
            checkNotNull(sourceVertex, "Could not find source vertex for media object: " + row.getLinkObjectId());

            String edgeId = getEdgeId(row);
            String edgeLabel = getEdgeLabel(row);
            getDataImporter().getGraph().addEdge(edgeId, sourceVertex, mediaVertex, edgeLabel, getDataImporter().getVisibility(), getDataImporter().getAuthorizations());
        } finally {
            in.close();
        }
    }

    private String getEdgeLabel(PtMediaAndValue row) {
        return getDataImporter().getHasMediaConceptTypeIri();
    }<commit_msg>choose when to decompress zlib better<commit_after>package io.lumify.palantir.dataImport;

import io.lumify.core.model.properties.LumifyProperties;
import io.lumify.palantir.dataImport.model.PtMediaAndValue;
import org.securegraph.Vertex;
import org.securegraph.VertexBuilder;
import org.securegraph.property.StreamingPropertyValue;

import java.io.BufferedInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.util.Arrays;
import java.util.zip.InflaterInputStream;

import static com.google.common.base.Preconditions.checkNotNull;

public class PtMediaAndValueImporter extends PtImporterBase<PtMediaAndValue> {
    private static final byte[] ZLIB_HEADER = new byte[2];

    static {
        ZLIB_HEADER[0] = (byte) 0x78;
        ZLIB_HEADER[1] = (byte) 0x9c;
    }

    public PtMediaAndValueImporter(DataImporter dataImporter) {
        super(dataImporter, PtMediaAndValue.class);
    }

    @Override
    protected void beforeProcessRows() {
        super.beforeProcessRows();
        getDataImporter().getGraph().flush();
    }

    @Override
    protected void processRow(PtMediaAndValue row) throws Exception {
        String propertyKey = getDataImporter().getIdPrefix() + row.getId();
        InputStream in = getContentInputStream(row);
        try {
            StreamingPropertyValue propertyValue = new StreamingPropertyValue(in, byte[].class);
            propertyValue.store(true);
            propertyValue.searchIndex(false);

            VertexBuilder vertexBuilder = getDataImporter().getGraph().prepareVertex(getMediaId(row), getDataImporter().getVisibility());
            LumifyProperties.RAW.addPropertyValue(vertexBuilder, propertyKey, propertyValue, getDataImporter().getVisibility());
            LumifyProperties.TITLE.setProperty(vertexBuilder, row.getTitle(), getDataImporter().getVisibility());
            Vertex mediaVertex = vertexBuilder.save(getDataImporter().getAuthorizations());

            Vertex sourceVertex = getVertexCache().get(getObjectId(row.getLinkObjectId()));
            checkNotNull(sourceVertex, "Could not find source vertex for media object: " + row.getLinkObjectId());

            String edgeId = getEdgeId(row);
            String edgeLabel = getEdgeLabel(row);
            getDataImporter().getGraph().addEdge(edgeId, sourceVertex, mediaVertex, edgeLabel, getDataImporter().getVisibility(), getDataImporter().getAuthorizations());
        } finally {
            in.close();
        }
    }

    private InputStream getContentInputStream(PtMediaAndValue row) throws IOException {
        byte[] header = new byte[2];
        InputStream in = new BufferedInputStream(row.getContents(), 100);

        in.mark(2);
        in.read(header);
        in.reset();

        if (Arrays.equals(header, ZLIB_HEADER)) {
            in = new InflaterInputStream(in);
        }
        return in;
    }

    private String getEdgeLabel(PtMediaAndValue row) {
        return getDataImporter().getHasMediaConceptTypeIri();
    }
{'fraction_non_alphanumeric': 0.10052724077328647, 'fraction_numerical': 0.0036906854130052723, 'mean_word_length': 3.680098684210526, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 6, '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': '3115670', 'n_tokens_mistral': 1567, 'n_tokens_neox': 1513, 'n_words': 290}