D-Net

package eu.dnetlib.functionality.index.solr.feed;

import java.io.StringReader;

import java.io.StringWriter;

import java.util.HashMap;

import java.util.Iterator;

import java.util.List;

import javax.xml.stream.*;

import javax.xml.stream.events.Namespace;

import javax.xml.stream.events.StartElement;

import javax.xml.stream.events.XMLEvent;

import com.google.common.collect.Lists;

import eu.dnetlib.functionality.index.solr.feed.ResultTransformer.Mode;

import org.apache.solr.common.SolrInputDocument;

import org.apache.solr.common.StringUtils;

/**

 * Optimized version of the document parser, drop in replacement of InputDocumentFactory.

 *

 * <p>

 * Faster because:

 * </p>

 * <ul>

 * <li>Doesn't create a DOM for the full document</li>

 * <li>Doesn't execute xpaths agains the DOM</li>

 * <li>Quickly serialize the 'result' element directly in a string.</li>

 * <li>Uses less memory: less pressure on GC and allows more threads to process this in parallel</li>

 * </ul>

 *

 * <p>

 * This class is fully reentrant and can be invoked in parallel.

 * </p>

 *

 * @author marko

 *

 */

public class StreamingInputDocumentFactory extends InputDocumentFactory {

	protected static final String DEFAULTDNETRESULT = "dnetResult";

	protected static final String TARGETFIELDS = "targetFields";

	protected static final String INDEX_RECORD_ID_ELEMENT = "indexRecordIdentifier";

	protected static final String ROOT_ELEMENT = "indexRecord";

	protected static final int MAX_FIELD_LENGTH = 32765;

	protected ThreadLocal<XMLInputFactory> inputFactory = new ThreadLocal<XMLInputFactory>() {

		@Override

		protected XMLInputFactory initialValue() {

			return XMLInputFactory.newInstance();

		}

	};

	protected ThreadLocal<XMLOutputFactory> outputFactory = new ThreadLocal<XMLOutputFactory>() {

		@Override

		protected XMLOutputFactory initialValue() {

			return XMLOutputFactory.newInstance();

		}

	};

	protected ThreadLocal<XMLEventFactory> eventFactory = new ThreadLocal<XMLEventFactory>() {

		@Override

		protected XMLEventFactory initialValue() {

			return XMLEventFactory.newInstance();

		}

	};

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.solr.feed.InputDocumentFactory#parseDocument(eu.dnetlib.functionality.index.solr.feed.IndexDocument,

	 * java.lang.String)

	 */

	@Override

	public SolrInputDocument parseDocument(final String version, final String inputDocument, final String dsId, final String resultName)

			throws XMLStreamException {

		return parseDocument(version, inputDocument, dsId, resultName, null);

	}

	/**

	 * {@inheritDoc}

	 *

	 * @see eu.dnetlib.functionality.index.solr.feed.InputDocumentFactory#parseDocument(eu.dnetlib.functionality.index.solr.feed.IndexDocument,

	 *      java.lang.String, com.google.common.base.Function)

	 */

	@Override

	public SolrInputDocument parseDocument(final String version,

			final String inputDocument,

			final String dsId,

			final String resultName,

			final ResultTransformer resultTransformer) {

		final StringWriter results = new StringWriter();

		final List<Namespace> nsList = Lists.newLinkedList();

		try {

			XMLEventReader parser = inputFactory.get().createXMLEventReader(new StringReader(inputDocument));

			final SolrInputDocument indexDocument = new SolrInputDocument(new HashMap<>());

			while (parser.hasNext()) {

				final XMLEvent event = parser.nextEvent();

				if ((event != null) && event.isStartElement()) {

					final String localName = event.asStartElement().getName().getLocalPart();

					if (ROOT_ELEMENT.equals(localName)) {

						nsList.addAll(getNamespaces(event));

					} else if (INDEX_RECORD_ID_ELEMENT.equals(localName)) {

						final XMLEvent text = parser.nextEvent();

						String recordId = getText(text);

						indexDocument.addField(INDEX_RECORD_ID, recordId);

					} else if (TARGETFIELDS.equals(localName)) {

						parseTargetFields(indexDocument, parser);

					} else if (resultName.equals(localName)) {

						if (resultTransformer == null || !(Mode.empty.equals(resultTransformer.getMode()))) {

							copyResult(indexDocument, results, parser, nsList, resultName, resultTransformer);

						}

					}

				}

			}

			if (version != null) {

				indexDocument.addField(DS_VERSION, version);

			}

			if (dsId != null) {

				indexDocument.addField(DS_ID, dsId);

			}

			if (!indexDocument.containsKey(INDEX_RECORD_ID)) {

				indexDocument.clear();

				System.err.println("missing indexrecord id:\n" + inputDocument);

			}

			return indexDocument;

		} catch (XMLStreamException e) {

			return new SolrInputDocument();

		}

	}

	private List<Namespace> getNamespaces(final XMLEvent event) {

		final List<Namespace> res = Lists.newLinkedList();

		@SuppressWarnings("unchecked")

		Iterator<Namespace> nsIter = event.asStartElement().getNamespaces();

		while (nsIter.hasNext()) {

			Namespace ns = nsIter.next();

			res.add(ns);

		}

		return res;

	}

	/**

	 * Parse the targetFields block and add fields to the solr document.

	 *

	 * @param indexDocument

	 * @param parser

	 * @throws XMLStreamException

	 */

	protected void parseTargetFields(final SolrInputDocument indexDocument, final XMLEventReader parser) throws XMLStreamException {

		boolean hasFields = false;

		while (parser.hasNext()) {

			final XMLEvent targetEvent = parser.nextEvent();

			if (targetEvent.isEndElement() && targetEvent.asEndElement().getName().getLocalPart().equals(TARGETFIELDS)) {

				break;

			}

			if (targetEvent.isStartElement()) {

				final String fieldName = targetEvent.asStartElement().getName().getLocalPart();

				final XMLEvent text = parser.nextEvent();

				String data = getText(text);

				addField(indexDocument, fieldName, data);

				hasFields = true;

			}

		}

		if (!hasFields) {

			indexDocument.clear();

		}

	}

	/**

	 * Copy the /indexRecord/result element and children, preserving namespace declarations etc.

	 *

	 * @param indexDocument

	 * @param results

	 * @param parser

	 * @param nsList

	 * @throws XMLStreamException

	 */

	protected void copyResult(final SolrInputDocument indexDocument,

			final StringWriter results,

			final XMLEventReader parser,

			final List<Namespace> nsList,

			final String dnetResult,

			final ResultTransformer resultTransformer) throws XMLStreamException {

		final XMLEventWriter writer = outputFactory.get().createXMLEventWriter(results);

		for (Namespace ns : nsList) {

			eventFactory.get().createNamespace(ns.getPrefix(), ns.getNamespaceURI());

		}

		StartElement newRecord = eventFactory.get().createStartElement("", null, RESULT, null, nsList.iterator());

		// new root record

		writer.add(newRecord);

		// copy the rest as it is

		while (parser.hasNext()) {

			final XMLEvent resultEvent = parser.nextEvent();

			// TODO: replace with depth tracking instead of close tag tracking.

			if (resultEvent.isEndElement() && resultEvent.asEndElement().getName().getLocalPart().equals(dnetResult)) {

				writer.add(eventFactory.get().createEndElement("", null, RESULT));

				break;

			}

			writer.add(resultEvent);

		}

		writer.close();

		if (resultTransformer != null) {

			indexDocument.addField(INDEX_RESULT, resultTransformer.apply(results.toString()));

		} else {

			indexDocument.addField(INDEX_RESULT, results.toString());

		}

	}

	/**

	 * Helper used to add a field to a solr doc. It avoids to add empy fields

	 *

	 * @param indexDocument

	 * @param field

	 * @param value

	 */

	private final void addField(final SolrInputDocument indexDocument, final String field, final String value) {

		String cleaned = value.trim();

		if (!cleaned.isEmpty()) {

			// log.info("\n\n adding field " + field.toLowerCase() + " value: " + cleaned + "\n");

			indexDocument.addField(field.toLowerCase(), cleaned);

		}

	}

	/**

	 * Helper used to get the string from a text element.

	 *

	 * @param text

	 * @return

	 */

	protected final String getText(final XMLEvent text) {

		if (text.isEndElement()) // log.warn("skipping because isEndOfElement " + text.asEndElement().getName().getLocalPart());

			return "";

		final String data = text.asCharacters().getData();

		if (data != null && data.length() > MAX_FIELD_LENGTH) {

			return data.substring(0, MAX_FIELD_LENGTH);

		}

		return data;

	}

}