unsafe/sort/UnsafeSortDataFormat.java

0001 /*
0002  * Licensed to the Apache Software Foundation (ASF) under one or more
0003  * contributor license agreements.  See the NOTICE file distributed with
0004  * this work for additional information regarding copyright ownership.
0005  * The ASF licenses this file to You under the Apache License, Version 2.0
0006  * (the "License"); you may not use this file except in compliance with
0007  * the License.  You may obtain a copy of the License at
0008  *
0009  *    http://www.apache.org/licenses/LICENSE-2.0
0010  *
0011  * Unless required by applicable law or agreed to in writing, software
0012  * distributed under the License is distributed on an "AS IS" BASIS,
0013  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
0014  * See the License for the specific language governing permissions and
0015  * limitations under the License.
0016  */
0017
0018 package org.apache.spark.util.collection.unsafe.sort;
0019
0020 import org.apache.spark.unsafe.Platform;
0021 import org.apache.spark.unsafe.array.LongArray;
0022 import org.apache.spark.util.collection.SortDataFormat;
0023
0024 /**
0025  * Supports sorting an array of (record pointer, key prefix) pairs.
0026  * Used in {@link UnsafeInMemorySorter}.
0027  * <p>
0028  * Within each long[] buffer, position {@code 2 * i} holds a pointer to the record at
0029  * index {@code i}, while position {@code 2 * i + 1} in the array holds an 8-byte key prefix.
0030  */
0031 public final class UnsafeSortDataFormat
0032   extends SortDataFormat<RecordPointerAndKeyPrefix, LongArray> {
0033
0034   private final LongArray buffer;
0035
0036   public UnsafeSortDataFormat(LongArray buffer) {
0037     this.buffer = buffer;
0038   }
0039
0040   @Override
0041   public RecordPointerAndKeyPrefix getKey(LongArray data, int pos) {
0042     // Since we re-use keys, this method shouldn't be called.
0043     throw new UnsupportedOperationException();
0044   }
0045
0046   @Override
0047   public RecordPointerAndKeyPrefix newKey() {
0048     return new RecordPointerAndKeyPrefix();
0049   }
0050
0051   @Override
0052   public RecordPointerAndKeyPrefix getKey(LongArray data, int pos,
0053                                           RecordPointerAndKeyPrefix reuse) {
0054     reuse.recordPointer = data.get(pos * 2);
0055     reuse.keyPrefix = data.get(pos * 2 + 1);
0056     return reuse;
0057   }
0058
0059   @Override
0060   public void swap(LongArray data, int pos0, int pos1) {
0061     long tempPointer = data.get(pos0 * 2);
0062     long tempKeyPrefix = data.get(pos0 * 2 + 1);
0063     data.set(pos0 * 2, data.get(pos1 * 2));
0064     data.set(pos0 * 2 + 1, data.get(pos1 * 2 + 1));
0065     data.set(pos1 * 2, tempPointer);
0066     data.set(pos1 * 2 + 1, tempKeyPrefix);
0067   }
0068
0069   @Override
0070   public void copyElement(LongArray src, int srcPos, LongArray dst, int dstPos) {
0071     dst.set(dstPos * 2, src.get(srcPos * 2));
0072     dst.set(dstPos * 2 + 1, src.get(srcPos * 2 + 1));
0073   }
0074
0075   @Override
0076   public void copyRange(LongArray src, int srcPos, LongArray dst, int dstPos, int length) {
0077     Platform.copyMemory(
0078       src.getBaseObject(),
0079       src.getBaseOffset() + srcPos * 16L,
0080       dst.getBaseObject(),
0081       dst.getBaseOffset() + dstPos * 16L,
0082       length * 16L);
0083   }
0084
0085   @Override
0086   public LongArray allocate(int length) {
0087     assert (length * 2L <= buffer.size()) :
0088       "the buffer is smaller than required: " + buffer.size() + " < " + (length * 2);
0089     return buffer;
0090   }
0091
0092 }