spark/unsafe/UTF8StringBuilder.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.unsafe;
0019
0020 import org.apache.spark.unsafe.array.ByteArrayMethods;
0021 import org.apache.spark.unsafe.types.UTF8String;
0022
0023 /**
0024  * A helper class to write {@link UTF8String}s to an internal buffer and build the concatenated
0025  * {@link UTF8String} at the end.
0026  */
0027 public class UTF8StringBuilder {
0028
0029   private static final int ARRAY_MAX = ByteArrayMethods.MAX_ROUNDED_ARRAY_LENGTH;
0030
0031   private byte[] buffer;
0032   private int cursor = Platform.BYTE_ARRAY_OFFSET;
0033
0034   public UTF8StringBuilder() {
0035     // Since initial buffer size is 16 in `StringBuilder`, we set the same size here
0036     this(16);
0037   }
0038
0039   public UTF8StringBuilder(int initialSize) {
0040     if (initialSize < 0) {
0041       throw new IllegalArgumentException("Size must be non-negative");
0042     }
0043     if (initialSize > ARRAY_MAX) {
0044       throw new IllegalArgumentException(
0045         "Size " + initialSize + " exceeded maximum size of " + ARRAY_MAX);
0046     }
0047     this.buffer = new byte[initialSize];
0048   }
0049
0050   // Grows the buffer by at least `neededSize`
0051   private void grow(int neededSize) {
0052     if (neededSize > ARRAY_MAX - totalSize()) {
0053       throw new UnsupportedOperationException(
0054         "Cannot grow internal buffer by size " + neededSize + " because the size after growing " +
0055           "exceeds size limitation " + ARRAY_MAX);
0056     }
0057     final int length = totalSize() + neededSize;
0058     if (buffer.length < length) {
0059       int newLength = length < ARRAY_MAX / 2 ? length * 2 : ARRAY_MAX;
0060       final byte[] tmp = new byte[newLength];
0061       Platform.copyMemory(
0062         buffer,
0063         Platform.BYTE_ARRAY_OFFSET,
0064         tmp,
0065         Platform.BYTE_ARRAY_OFFSET,
0066         totalSize());
0067       buffer = tmp;
0068     }
0069   }
0070
0071   private int totalSize() {
0072     return cursor - Platform.BYTE_ARRAY_OFFSET;
0073   }
0074
0075   public void append(UTF8String value) {
0076     grow(value.numBytes());
0077     value.writeToMemory(buffer, cursor);
0078     cursor += value.numBytes();
0079   }
0080
0081   public void append(String value) {
0082     append(UTF8String.fromString(value));
0083   }
0084
0085   public void appendBytes(Object base, long offset, int length) {
0086     grow(length);
0087     Platform.copyMemory(
0088       base,
0089       offset,
0090       buffer,
0091       cursor,
0092       length);
0093     cursor += length;
0094   }
0095
0096   public UTF8String build() {
0097     return UTF8String.fromBytes(buffer, 0, totalSize());
0098   }
0099 }