/*
* Hibernate, Relational Persistence for Idiomatic Java
*
* JBoss, Home of Professional Open Source
* Copyright 2011 Red Hat Inc. and/or its affiliates and other contributors
* as indicated by the @authors tag. All rights reserved.
* See the copyright.txt in the distribution for a
* full listing of individual contributors.
*
* This copyrighted material is made available to anyone wishing to use,
* modify, copy, or redistribute it subject to the terms and conditions
* of the GNU Lesser General Public License, v. 2.1.
* This program is distributed in the hope that it will be useful, but WITHOUT A
* WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
* PARTICULAR PURPOSE. See the GNU Lesser General Public License for more details.
* You should have received a copy of the GNU Lesser General Public License,
* v.2.1 along with this distribution; if not, write to the Free Software
* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
* MA 02110-1301, USA.
*/
package org.hibernate.search.indexes.serialization.avro.impl;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.Serializable;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import org.apache.avro.Protocol;
import org.apache.avro.Schema;
import org.apache.avro.generic.GenericData;
import org.apache.avro.generic.GenericDatumWriter;
import org.apache.avro.generic.GenericRecord;
import org.apache.avro.io.BinaryEncoder;
import org.apache.avro.io.EncoderFactory;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.CharTermAttributeImpl;
import org.apache.lucene.analysis.tokenattributes.FlagsAttribute;
import org.apache.lucene.analysis.tokenattributes.KeywordAttribute;
import org.apache.lucene.analysis.tokenattributes.OffsetAttribute;
import org.apache.lucene.analysis.tokenattributes.PayloadAttribute;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.apache.lucene.analysis.tokenattributes.TypeAttribute;
import org.apache.lucene.document.Fieldable;
import org.apache.lucene.util.AttributeImpl;
import org.apache.solr.handler.AnalysisRequestHandlerBase;
import org.hibernate.search.backend.LuceneWork;
import org.hibernate.search.indexes.serialization.spi.LuceneFieldContext;
import org.hibernate.search.indexes.serialization.spi.LuceneNumericFieldContext;
import org.hibernate.search.indexes.serialization.spi.Serializer;
import org.hibernate.search.util.logging.impl.Log;
import org.hibernate.search.util.logging.impl.LoggerFactory;
import static org.hibernate.search.indexes.serialization.impl.SerializationHelper.toByteArray;
/**
* @author Emmanuel Bernard <emmanuel@hibernate.org>
*/
public class AvroSerializer implements Serializer {
private static final Log log = LoggerFactory.make();
private GenericRecord idRecord;
private List<GenericRecord> fieldables;
private List<GenericRecord> operations;
private List<String> classReferences;
private GenericRecord document;
private final Protocol protocol;
public AvroSerializer(Protocol protocol) {
this.protocol = protocol;
this.classReferences = new ArrayList<String>();
}
@Override
public void luceneWorks(List<LuceneWork> works) {
operations = new ArrayList<GenericRecord>( works.size() );
}
@Override
public void addOptimizeAll() {
operations.add( new GenericData.Record( protocol.getType( "OptimizeAll" ) ) );
}
@Override
public void addPurgeAll(String entityClassName) {
int classRef = getClassReference( entityClassName );
GenericRecord purgeAll = new GenericData.Record( protocol.getType( "PurgeAll" ) );
purgeAll.put( "class", classRef );
operations.add( purgeAll );
}
private int getClassReference(String entityClassName) {
int classRef = classReferences.indexOf( entityClassName );
if ( classRef == -1 ) {
classReferences.add( entityClassName );
classRef = classReferences.size() - 1;
}
return classRef;
}
@Override
public void addIdSerializedInJava(byte[] id) {
this.idRecord = new GenericData.Record( protocol.getType( "Id" ) );
idRecord.put( "value", ByteBuffer.wrap( id ) );
}
@Override
public void addIdAsInteger(int id) {
this.idRecord = new GenericData.Record( protocol.getType( "Id" ) );
idRecord.put( "value", id );
}
@Override
public void addIdAsLong(long id) {
this.idRecord = new GenericData.Record( protocol.getType( "Id" ) );
idRecord.put( "value", id );
}
@Override
public void addIdAsFloat(float id) {
this.idRecord = new GenericData.Record( protocol.getType( "Id" ) );
idRecord.put( "value", id );
}
@Override
public void addIdAsDouble(double id) {
this.idRecord = new GenericData.Record( protocol.getType( "Id" ) );
idRecord.put( "value", id );
}
@Override
public void addIdAsString(String id) {
this.idRecord = new GenericData.Record( protocol.getType( "Id" ) );
idRecord.put( "value", id );
}
@Override
public void addDelete(String entityClassName) {
int classRef = getClassReference( entityClassName );
GenericRecord delete = new GenericData.Record( protocol.getType( "Delete" ) );
delete.put( "class", classRef );
delete.put( "id", idRecord );
operations.add( delete );
idRecord = null;
}
@Override
public void addAdd(String entityClassName, Map<String, String> fieldToAnalyzerMap) {
int classRef = getClassReference( entityClassName );
GenericRecord add = new GenericData.Record( protocol.getType( "Add" ) );
add.put( "class", classRef );
add.put( "id", idRecord );
add.put( "document", document );
add.put( "fieldToAnalyzerMap", fieldToAnalyzerMap );
operations.add( add );
idRecord = null;
clearDocument();
}
@Override
public void addUpdate(String entityClassName, Map<String, String> fieldToAnalyzerMap) {
int classRef = getClassReference( entityClassName );
GenericRecord update = new GenericData.Record( protocol.getType( "Update" ) );
update.put( "class", classRef );
update.put( "id", idRecord );
update.put( "document", document );
update.put( "fieldToAnalyzerMap", fieldToAnalyzerMap );
operations.add( update );
idRecord = null;
clearDocument();
}
@Override
public byte[] serialize() {
final ByteArrayOutputStream out = new ByteArrayOutputStream();
out.write( AvroSerializationProvider.getMajorVersion() );
out.write( AvroSerializationProvider.getMinorVersion() );
Schema msgSchema = protocol.getType( "Message" );
GenericDatumWriter<GenericRecord> writer = new GenericDatumWriter<GenericRecord>( msgSchema );
BinaryEncoder encoder = EncoderFactory.get().directBinaryEncoder( out, null );
GenericRecord message = new GenericData.Record( msgSchema );
message.put( "classReferences", classReferences );
message.put( "operations", operations );
operations = null;
try {
writer.write( message, encoder );
encoder.flush();
}
catch ( IOException e ) {
throw log.unableToSerializeInAvro(e);
}
return out.toByteArray();
}
@Override
public void fields(List<Fieldable> fields) {
fieldables = new ArrayList<GenericRecord>( fields.size() );
}
@Override
public void addIntNumericField(int value, LuceneNumericFieldContext context) {
GenericRecord numericField = createNumericfield( "NumericIntField", context );
numericField.put( "value", value );
fieldables.add( numericField );
}
private GenericRecord createNumericfield(String schemaName, LuceneNumericFieldContext context) {
GenericRecord numericField = new GenericData.Record( protocol.getType( schemaName ) );
numericField.put( "name", context.getName() );
numericField.put( "precisionStep", context.getPrecisionStep() );
numericField.put( "store", context.getStore() );
numericField.put( "indexed", context.isIndexed() );
numericField.put( "boost", context.getBoost() );
numericField.put( "omitNorms", context.getOmitNorms() );
numericField.put( "omitTermFreqAndPositions", context.getOmitTermFreqAndPositions() );
return numericField;
}
@Override
public void addLongNumericField(long value, LuceneNumericFieldContext context) {
GenericRecord numericField = createNumericfield( "NumericLongField", context );
numericField.put( "value", value );
fieldables.add( numericField );
}
@Override
public void addFloatNumericField(float value, LuceneNumericFieldContext context) {
GenericRecord numericField = createNumericfield( "NumericFloatField", context );
numericField.put( "value", value );
fieldables.add( numericField );
}
@Override
public void addDoubleNumericField(double value, LuceneNumericFieldContext context) {
GenericRecord numericField = createNumericfield( "NumericDoubleField", context );
numericField.put( "value", value );
fieldables.add( numericField );
}
@Override
public void addFieldWithBinaryData(LuceneFieldContext context) {
GenericRecord field = createNormalField( "BinaryField", context );
field.put( "offset", context.getBinaryOffset() );
field.put( "length", context.getBinaryLength() );
field.put( "value", ByteBuffer.wrap( context.getBinaryValue() ) );
fieldables.add( field );
}
private GenericRecord createNormalField(String schemaName, LuceneFieldContext context) {
GenericRecord field = new GenericData.Record( protocol.getType( schemaName ) );
field.put( "name", context.getName() );
field.put( "boost", context.getBoost() );
field.put( "omitNorms", context.isOmitNorms() );
field.put( "omitTermFreqAndPositions", context.isOmitTermFreqAndPositions() );
return field;
}
@Override
public void addFieldWithStringData(LuceneFieldContext context) {
GenericRecord field = createNormalField( "StringField", context );
field.put( "value", context.getStringValue() );
field.put( "store", context.getStore() );
field.put( "index", context.getIndex() );
field.put( "termVector", context.getTermVector() );
fieldables.add( field );
}
@Override
public void addFieldWithTokenStreamData(LuceneFieldContext context) {
GenericRecord field = createNormalField( "TokenStreamField", context );
List<List<AttributeImpl>> stream = context.getTokenStream().getStream();
List<List<Object>> value = new ArrayList<List<Object>>( stream.size() );
for( List<AttributeImpl> attrs : stream ) {
List<Object> elements = new ArrayList<Object>( attrs.size() );
for(AttributeImpl attr : attrs) {
elements.add( buildAttributeImpl( attr ) );
}
value.add(elements);
}
field.put( "value", value );
field.put( "termVector", context.getTermVector() );
fieldables.add( field );
}
private Object buildAttributeImpl(AttributeImpl attr) {
if ( attr instanceof AnalysisRequestHandlerBase.TokenTrackingAttributeImpl ) {
GenericRecord record = new GenericData.Record( protocol.getType( "TokenTrackingAttribute" ) );
int[] positions = ( (AnalysisRequestHandlerBase.TokenTrackingAttributeImpl) attr ).getPositions();
List<Integer> fullPositions = new ArrayList<Integer>( positions.length );
for (int position : positions) {
fullPositions.add( position );
}
record.put( "positions", fullPositions );
return record;
}
else if (attr instanceof CharTermAttributeImpl) {
GenericRecord record = new GenericData.Record( protocol.getType( "CharTermAttribute" ) );
CharTermAttribute charAttr = (CharTermAttribute) attr;
record.put("sequence", charAttr.toString() );
return record;
}
else if (attr instanceof PayloadAttribute) {
GenericRecord record = new GenericData.Record( protocol.getType( "PayloadAttribute" ) );
PayloadAttribute payloadAttr = (PayloadAttribute) attr;
record.put("payload", ByteBuffer.wrap( payloadAttr.getPayload().toByteArray() ) );
return record;
}
else if (attr instanceof KeywordAttribute) {
GenericRecord record = new GenericData.Record( protocol.getType( "KeywordAttribute" ) );
KeywordAttribute narrowedAttr = (KeywordAttribute) attr;
record.put("isKeyword", narrowedAttr.isKeyword() );
return record;
}
else if (attr instanceof PositionIncrementAttribute ) {
GenericRecord record = new GenericData.Record( protocol.getType( "PositionIncrementAttribute" ) );
PositionIncrementAttribute narrowedAttr = (PositionIncrementAttribute) attr;
record.put("positionIncrement", narrowedAttr.getPositionIncrement() );
return record;
}
else if (attr instanceof FlagsAttribute ) {
GenericRecord record = new GenericData.Record( protocol.getType( "FlagsAttribute" ) );
FlagsAttribute narrowedAttr = (FlagsAttribute) attr;
record.put("flags", narrowedAttr.getFlags() );
return record;
}
else if (attr instanceof TypeAttribute ) {
GenericRecord record = new GenericData.Record( protocol.getType( "TypeAttribute" ) );
TypeAttribute narrowedAttr = (TypeAttribute) attr;
record.put("type", narrowedAttr.type() );
return record;
}
else if (attr instanceof OffsetAttribute ) {
GenericRecord record = new GenericData.Record( protocol.getType( "OffsetAttribute" ) );
OffsetAttribute narrowedAttr = (OffsetAttribute) attr;
record.put("startOffset", narrowedAttr.startOffset() );
record.put("endOffset", narrowedAttr.endOffset() );
return record;
}
else if (attr instanceof Serializable) {
return ByteBuffer.wrap( toByteArray(attr) );
}
else {
throw log.attributeNotRecognizedNorSerializable( attr.getClass() );
}
}
@Override
public void addFieldWithSerializableReaderData(LuceneFieldContext context) {
GenericRecord field = createNormalField( "ReaderField", context );
field.put( "value", ByteBuffer.wrap( context.getReaderValue() ) );
field.put( "termVector", context.getTermVector() );
fieldables.add( field );
}
@Override
public void addFieldWithSerializableFieldable(byte[] fieldable) {
GenericRecord customFieldable = new GenericData.Record( protocol.getType( "CustomFieldable" ) );
customFieldable.put( "instance", ByteBuffer.wrap( fieldable ) );
fieldables.add( customFieldable );
}
@Override
public void addDocument(float boost) {
document = new GenericData.Record( protocol.getType( "Document" ) );
document.put( "boost", boost );
document.put( "fieldables", fieldables );
}
private void clearDocument() {
document = null;
fieldables = null;
}
}