lucene/core/src/java/org/apache/lucene/index/SortedSetDocValuesWriter.java - lucene-solr - Git at Google

 package org.apache.lucene.index;

 /*
  * Licensed to the Apache Software Foundation (ASF) under one or more
  * contributor license agreements.  See the NOTICE file distributed with
  * this work for additional information regarding copyright ownership.
  * The ASF licenses this file to You under the Apache License, Version 2.0
  * (the "License"); you may not use this file except in compliance with
  * the License.  You may obtain a copy of the License at
  *
  *     http://www.apache.org/licenses/LICENSE-2.0
  *
  * Unless required by applicable law or agreed to in writing, software
  * distributed under the License is distributed on an "AS IS" BASIS,
  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */

 import static org.apache.lucene.util.ByteBlockPool.BYTE_BLOCK_SIZE;

 import java.io.IOException;
 import java.util.Arrays;
 import java.util.Iterator;
 import java.util.NoSuchElementException;

 import org.apache.lucene.codecs.DocValuesConsumer;
 import org.apache.lucene.util.ArrayUtil;
 import org.apache.lucene.util.ByteBlockPool;
 import org.apache.lucene.util.BytesRef;
 import org.apache.lucene.util.BytesRefHash.DirectBytesStartArray;
 import org.apache.lucene.util.BytesRefHash;
 import org.apache.lucene.util.Counter;
 import org.apache.lucene.util.RamUsageEstimator;
 import org.apache.lucene.util.packed.AppendingDeltaPackedLongBuffer;
 import org.apache.lucene.util.packed.AppendingPackedLongBuffer;
 import org.apache.lucene.util.packed.PackedInts;

 /** Buffers up pending byte[]s per doc, deref and sorting via
  *  int ord, then flushes when segment flushes. */
 class SortedSetDocValuesWriter extends DocValuesWriter {
   final BytesRefHash hash;
   private AppendingPackedLongBuffer pending; // stream of all termIDs
   private AppendingDeltaPackedLongBuffer pendingCounts; // termIDs per doc
   private final Counter iwBytesUsed;
   private long bytesUsed; // this only tracks differences in 'pending' and 'pendingCounts'
   private final FieldInfo fieldInfo;
   private int currentDoc;
   private int currentValues[] = new int[8];
   private int currentUpto = 0;
   private int maxCount = 0;

   public SortedSetDocValuesWriter(FieldInfo fieldInfo, Counter iwBytesUsed) {
     this.fieldInfo = fieldInfo;
     this.iwBytesUsed = iwBytesUsed;
     hash = new BytesRefHash(
         new ByteBlockPool(
             new ByteBlockPool.DirectTrackingAllocator(iwBytesUsed)),
             BytesRefHash.DEFAULT_CAPACITY,
             new DirectBytesStartArray(BytesRefHash.DEFAULT_CAPACITY, iwBytesUsed));
     pending = new AppendingPackedLongBuffer(PackedInts.COMPACT);
     pendingCounts = new AppendingDeltaPackedLongBuffer(PackedInts.COMPACT);
     bytesUsed = pending.ramBytesUsed() + pendingCounts.ramBytesUsed();
     iwBytesUsed.addAndGet(bytesUsed);
   }

   public void addValue(int docID, BytesRef value) {
     if (value == null) {
       throw new IllegalArgumentException("field \"" + fieldInfo.name + "\": null value not allowed");
     }
     if (value.length > (BYTE_BLOCK_SIZE - 2)) {
       throw new IllegalArgumentException("DocValuesField \"" + fieldInfo.name + "\" is too large, must be <= " + (BYTE_BLOCK_SIZE - 2));
     }

     if (docID != currentDoc) {
       finishCurrentDoc();
     }

     // Fill in any holes:
     while(currentDoc < docID) {
       pendingCounts.add(0); // no values
       currentDoc++;
     }

     addOneValue(value);
     updateBytesUsed();
   }

   // finalize currentDoc: this deduplicates the current term ids
   private void finishCurrentDoc() {
     Arrays.sort(currentValues, 0, currentUpto);
     int lastValue = -1;
     int count = 0;
     for (int i = 0; i < currentUpto; i++) {
       int termID = currentValues[i];
       // if its not a duplicate
       if (termID != lastValue) {
         pending.add(termID); // record the term id
         count++;
       }
       lastValue = termID;
     }
     // record the number of unique term ids for this doc
     pendingCounts.add(count);
     maxCount = Math.max(maxCount, count);
     currentUpto = 0;
     currentDoc++;
   }

   @Override
   public void finish(int maxDoc) {
     finishCurrentDoc();

     // fill in any holes
     for (int i = currentDoc; i < maxDoc; i++) {
       pendingCounts.add(0); // no values
     }
   }

   private void addOneValue(BytesRef value) {
     int termID = hash.add(value);
     if (termID < 0) {
       termID = -termID-1;
     } else {
       // reserve additional space for each unique value:
       // 1. when indexing, when hash is 50% full, rehash() suddenly needs 2*size ints.
       //    TODO: can this same OOM happen in THPF?
       // 2. when flushing, we need 1 int per value (slot in the ordMap).
       iwBytesUsed.addAndGet(2 * RamUsageEstimator.NUM_BYTES_INT);
     }

     if (currentUpto == currentValues.length) {
       currentValues = ArrayUtil.grow(currentValues, currentValues.length+1);
       // reserve additional space for max # values per-doc
       // when flushing, we need an int[] to sort the mapped-ords within the doc
       iwBytesUsed.addAndGet((currentValues.length - currentUpto) * 2 * RamUsageEstimator.NUM_BYTES_INT);
     }

     currentValues[currentUpto] = termID;
     currentUpto++;
   }

   private void updateBytesUsed() {
     final long newBytesUsed = pending.ramBytesUsed() + pendingCounts.ramBytesUsed();
     iwBytesUsed.addAndGet(newBytesUsed - bytesUsed);
     bytesUsed = newBytesUsed;
   }

   @Override
   public void flush(SegmentWriteState state, DocValuesConsumer dvConsumer) throws IOException {
     final int maxDoc = state.segmentInfo.getDocCount();
     final int maxCountPerDoc = maxCount;
     assert pendingCounts.size() == maxDoc;
     final int valueCount = hash.size();

     final int[] sortedValues = hash.sort(BytesRef.getUTF8SortedAsUnicodeComparator());
     final int[] ordMap = new int[valueCount];

     for(int ord=0;ord<valueCount;ord++) {
       ordMap[sortedValues[ord]] = ord;
     }

     dvConsumer.addSortedSetField(fieldInfo,

                               // ord -> value
                               new Iterable<BytesRef>() {
                                 @Override
                                 public Iterator<BytesRef> iterator() {
                                   return new ValuesIterator(sortedValues, valueCount);
                                 }
                               },

                               // doc -> ordCount
                               new Iterable<Number>() {
                                 @Override
                                 public Iterator<Number> iterator() {
                                   return new OrdCountIterator(maxDoc);
                                 }
                               },

                               // ords
                               new Iterable<Number>() {
                                 @Override
                                 public Iterator<Number> iterator() {
                                   return new OrdsIterator(ordMap, maxCountPerDoc);
                                 }
                               });
   }

   @Override
   public void abort() {
   }

   // iterates over the unique values we have in ram
   private class ValuesIterator implements Iterator<BytesRef> {
     final int sortedValues[];
     final BytesRef scratch = new BytesRef();
     final int valueCount;
     int ordUpto;

     ValuesIterator(int sortedValues[], int valueCount) {
       this.sortedValues = sortedValues;
       this.valueCount = valueCount;
     }

     @Override
     public boolean hasNext() {
       return ordUpto < valueCount;
     }

     @Override
     public BytesRef next() {
       if (!hasNext()) {
         throw new NoSuchElementException();
       }
       hash.get(sortedValues[ordUpto], scratch);
       ordUpto++;
       return scratch;
     }

     @Override
     public void remove() {
       throw new UnsupportedOperationException();
     }
   }

   // iterates over the ords for each doc we have in ram
   private class OrdsIterator implements Iterator<Number> {
     final AppendingPackedLongBuffer.Iterator iter = pending.iterator();
     final AppendingDeltaPackedLongBuffer.Iterator counts = pendingCounts.iterator();
     final int ordMap[];
     final long numOrds;
     long ordUpto;

     final int currentDoc[];
     int currentUpto;
     int currentLength;

     OrdsIterator(int ordMap[], int maxCount) {
       this.currentDoc = new int[maxCount];
       this.ordMap = ordMap;
       this.numOrds = pending.size();
     }

     @Override
     public boolean hasNext() {
       return ordUpto < numOrds;
     }

     @Override
     public Number next() {
       if (!hasNext()) {
         throw new NoSuchElementException();
       }
       while (currentUpto == currentLength) {
         // refill next doc, and sort remapped ords within the doc.
         currentUpto = 0;
         currentLength = (int) counts.next();
         for (int i = 0; i < currentLength; i++) {
           currentDoc[i] = ordMap[(int) iter.next()];
         }
         Arrays.sort(currentDoc, 0, currentLength);
       }
       int ord = currentDoc[currentUpto];
       currentUpto++;
       ordUpto++;
       // TODO: make reusable Number
       return ord;
     }

     @Override
     public void remove() {
       throw new UnsupportedOperationException();
     }
   }

   private class OrdCountIterator implements Iterator<Number> {
     final AppendingDeltaPackedLongBuffer.Iterator iter = pendingCounts.iterator();
     final int maxDoc;
     int docUpto;

     OrdCountIterator(int maxDoc) {
       this.maxDoc = maxDoc;
       assert pendingCounts.size() == maxDoc;
     }

     @Override
     public boolean hasNext() {
       return docUpto < maxDoc;
     }

     @Override
     public Number next() {
       if (!hasNext()) {
         throw new NoSuchElementException();
       }
       docUpto++;
       // TODO: make reusable Number
       return iter.next();
     }

     @Override
     public void remove() {
       throw new UnsupportedOperationException();
     }
   }
 }
	package org.apache.lucene.index;

	/*
	* Licensed to the Apache Software Foundation (ASF) under one or more
	* contributor license agreements. See the NOTICE file distributed with
	* this work for additional information regarding copyright ownership.
	* The ASF licenses this file to You under the Apache License, Version 2.0
	* (the "License"); you may not use this file except in compliance with
	* the License. You may obtain a copy of the License at
	*
	* http://www.apache.org/licenses/LICENSE-2.0
	*
	* Unless required by applicable law or agreed to in writing, software
	* distributed under the License is distributed on an "AS IS" BASIS,
	* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
	* See the License for the specific language governing permissions and
	* limitations under the License.
	*/

	import static org.apache.lucene.util.ByteBlockPool.BYTE_BLOCK_SIZE;

	import java.io.IOException;
	import java.util.Arrays;
	import java.util.Iterator;
	import java.util.NoSuchElementException;

	import org.apache.lucene.codecs.DocValuesConsumer;
	import org.apache.lucene.util.ArrayUtil;
	import org.apache.lucene.util.ByteBlockPool;
	import org.apache.lucene.util.BytesRef;
	import org.apache.lucene.util.BytesRefHash.DirectBytesStartArray;
	import org.apache.lucene.util.BytesRefHash;
	import org.apache.lucene.util.Counter;
	import org.apache.lucene.util.RamUsageEstimator;
	import org.apache.lucene.util.packed.AppendingDeltaPackedLongBuffer;
	import org.apache.lucene.util.packed.AppendingPackedLongBuffer;
	import org.apache.lucene.util.packed.PackedInts;

	/** Buffers up pending byte[]s per doc, deref and sorting via
	* int ord, then flushes when segment flushes. */
	class SortedSetDocValuesWriter extends DocValuesWriter {
	final BytesRefHash hash;
	private AppendingPackedLongBuffer pending; // stream of all termIDs
	private AppendingDeltaPackedLongBuffer pendingCounts; // termIDs per doc
	private final Counter iwBytesUsed;
	private long bytesUsed; // this only tracks differences in 'pending' and 'pendingCounts'
	private final FieldInfo fieldInfo;
	private int currentDoc;
	private int currentValues[] = new int[8];
	private int currentUpto = 0;
	private int maxCount = 0;

	public SortedSetDocValuesWriter(FieldInfo fieldInfo, Counter iwBytesUsed) {
	this.fieldInfo = fieldInfo;
	this.iwBytesUsed = iwBytesUsed;
	hash = new BytesRefHash(
	new ByteBlockPool(
	new ByteBlockPool.DirectTrackingAllocator(iwBytesUsed)),
	BytesRefHash.DEFAULT_CAPACITY,
	new DirectBytesStartArray(BytesRefHash.DEFAULT_CAPACITY, iwBytesUsed));
	pending = new AppendingPackedLongBuffer(PackedInts.COMPACT);
	pendingCounts = new AppendingDeltaPackedLongBuffer(PackedInts.COMPACT);
	bytesUsed = pending.ramBytesUsed() + pendingCounts.ramBytesUsed();
	iwBytesUsed.addAndGet(bytesUsed);
	}

	public void addValue(int docID, BytesRef value) {
	if (value == null) {
	throw new IllegalArgumentException("field \"" + fieldInfo.name + "\": null value not allowed");
	}
	if (value.length > (BYTE_BLOCK_SIZE - 2)) {
	throw new IllegalArgumentException("DocValuesField \"" + fieldInfo.name + "\" is too large, must be <= " + (BYTE_BLOCK_SIZE - 2));
	}

	if (docID != currentDoc) {
	finishCurrentDoc();
	}

	// Fill in any holes:
	while(currentDoc < docID) {
	pendingCounts.add(0); // no values
	currentDoc++;
	}

	addOneValue(value);
	updateBytesUsed();
	}

	// finalize currentDoc: this deduplicates the current term ids
	private void finishCurrentDoc() {
	Arrays.sort(currentValues, 0, currentUpto);
	int lastValue = -1;
	int count = 0;
	for (int i = 0; i < currentUpto; i++) {
	int termID = currentValues[i];
	// if its not a duplicate
	if (termID != lastValue) {
	pending.add(termID); // record the term id
	count++;
	}
	lastValue = termID;
	}
	// record the number of unique term ids for this doc
	pendingCounts.add(count);
	maxCount = Math.max(maxCount, count);
	currentUpto = 0;
	currentDoc++;
	}

	@Override
	public void finish(int maxDoc) {
	finishCurrentDoc();

	// fill in any holes
	for (int i = currentDoc; i < maxDoc; i++) {
	pendingCounts.add(0); // no values
	}
	}

	private void addOneValue(BytesRef value) {
	int termID = hash.add(value);
	if (termID < 0) {
	termID = -termID-1;
	} else {
	// reserve additional space for each unique value:
	// 1. when indexing, when hash is 50% full, rehash() suddenly needs 2*size ints.
	// TODO: can this same OOM happen in THPF?
	// 2. when flushing, we need 1 int per value (slot in the ordMap).
	iwBytesUsed.addAndGet(2 * RamUsageEstimator.NUM_BYTES_INT);
	}

	if (currentUpto == currentValues.length) {
	currentValues = ArrayUtil.grow(currentValues, currentValues.length+1);
	// reserve additional space for max # values per-doc
	// when flushing, we need an int[] to sort the mapped-ords within the doc
	iwBytesUsed.addAndGet((currentValues.length - currentUpto) * 2 * RamUsageEstimator.NUM_BYTES_INT);
	}

	currentValues[currentUpto] = termID;
	currentUpto++;
	}

	private void updateBytesUsed() {
	final long newBytesUsed = pending.ramBytesUsed() + pendingCounts.ramBytesUsed();
	iwBytesUsed.addAndGet(newBytesUsed - bytesUsed);
	bytesUsed = newBytesUsed;
	}

	@Override
	public void flush(SegmentWriteState state, DocValuesConsumer dvConsumer) throws IOException {
	final int maxDoc = state.segmentInfo.getDocCount();
	final int maxCountPerDoc = maxCount;
	assert pendingCounts.size() == maxDoc;
	final int valueCount = hash.size();

	final int[] sortedValues = hash.sort(BytesRef.getUTF8SortedAsUnicodeComparator());
	final int[] ordMap = new int[valueCount];

	for(int ord=0;ord<valueCount;ord++) {
	ordMap[sortedValues[ord]] = ord;
	}

	dvConsumer.addSortedSetField(fieldInfo,

	// ord -> value
	new Iterable<BytesRef>() {
	@Override
	public Iterator<BytesRef> iterator() {
	return new ValuesIterator(sortedValues, valueCount);
	}
	},

	// doc -> ordCount
	new Iterable<Number>() {
	@Override
	public Iterator<Number> iterator() {
	return new OrdCountIterator(maxDoc);
	}
	},

	// ords
	new Iterable<Number>() {
	@Override
	public Iterator<Number> iterator() {
	return new OrdsIterator(ordMap, maxCountPerDoc);
	}
	});
	}

	@Override
	public void abort() {
	}

	// iterates over the unique values we have in ram
	private class ValuesIterator implements Iterator<BytesRef> {
	final int sortedValues[];
	final BytesRef scratch = new BytesRef();
	final int valueCount;
	int ordUpto;

	ValuesIterator(int sortedValues[], int valueCount) {
	this.sortedValues = sortedValues;
	this.valueCount = valueCount;
	}

	@Override
	public boolean hasNext() {
	return ordUpto < valueCount;
	}

	@Override
	public BytesRef next() {
	if (!hasNext()) {
	throw new NoSuchElementException();
	}
	hash.get(sortedValues[ordUpto], scratch);
	ordUpto++;
	return scratch;
	}

	@Override
	public void remove() {
	throw new UnsupportedOperationException();
	}
	}

	// iterates over the ords for each doc we have in ram
	private class OrdsIterator implements Iterator<Number> {
	final AppendingPackedLongBuffer.Iterator iter = pending.iterator();
	final AppendingDeltaPackedLongBuffer.Iterator counts = pendingCounts.iterator();
	final int ordMap[];
	final long numOrds;
	long ordUpto;

	final int currentDoc[];
	int currentUpto;
	int currentLength;

	OrdsIterator(int ordMap[], int maxCount) {
	this.currentDoc = new int[maxCount];
	this.ordMap = ordMap;
	this.numOrds = pending.size();
	}

	@Override
	public boolean hasNext() {
	return ordUpto < numOrds;
	}

	@Override
	public Number next() {
	if (!hasNext()) {
	throw new NoSuchElementException();
	}
	while (currentUpto == currentLength) {
	// refill next doc, and sort remapped ords within the doc.
	currentUpto = 0;
	currentLength = (int) counts.next();
	for (int i = 0; i < currentLength; i++) {
	currentDoc[i] = ordMap[(int) iter.next()];
	}
	Arrays.sort(currentDoc, 0, currentLength);
	}
	int ord = currentDoc[currentUpto];
	currentUpto++;
	ordUpto++;
	// TODO: make reusable Number
	return ord;
	}

	@Override
	public void remove() {
	throw new UnsupportedOperationException();
	}
	}

	private class OrdCountIterator implements Iterator<Number> {
	final AppendingDeltaPackedLongBuffer.Iterator iter = pendingCounts.iterator();
	final int maxDoc;
	int docUpto;

	OrdCountIterator(int maxDoc) {
	this.maxDoc = maxDoc;
	assert pendingCounts.size() == maxDoc;
	}

	@Override
	public boolean hasNext() {
	return docUpto < maxDoc;
	}

	@Override
	public Number next() {
	if (!hasNext()) {
	throw new NoSuchElementException();
	}
	docUpto++;
	// TODO: make reusable Number
	return iter.next();
	}

	@Override
	public void remove() {
	throw new UnsupportedOperationException();
	}
	}
	}