com.facebook.presto.spark.classloader_interface.PrestoSparkShuffleSerializer Maven / Gradle / Ivy
/*
* Licensed 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.
*/
package com.facebook.presto.spark.classloader_interface;
import org.apache.spark.serializer.DeserializationStream;
import org.apache.spark.serializer.SerializationStream;
import org.apache.spark.serializer.Serializer;
import org.apache.spark.serializer.SerializerInstance;
import scala.Tuple2;
import scala.collection.AbstractIterator;
import scala.collection.Iterator;
import scala.reflect.ClassTag;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.io.Serializable;
import java.io.UncheckedIOException;
import java.nio.ByteBuffer;
import java.util.NoSuchElementException;
import static java.util.Objects.requireNonNull;
public class PrestoSparkShuffleSerializer
extends Serializer
implements Serializable
{
@Override
public SerializerInstance newInstance()
{
return new PrestoSparkShuffleSerializerInstance();
}
@Override
public boolean supportsRelocationOfSerializedObjects()
{
return true;
}
public static class PrestoSparkShuffleSerializerInstance
extends SerializerInstance
{
private final PrestoSparkMutableRow row = new PrestoSparkMutableRow();
private final Tuple2 tuple = new Tuple2<>(null, row);
@Override
public SerializationStream serializeStream(OutputStream outputStream)
{
return new PrestoSparkShuffleSerializationStream(outputStream);
}
@Override
public DeserializationStream deserializeStream(InputStream inputStream)
{
return new PrestoSparkShuffleDeserializationStream(inputStream);
}
@Override
public ByteBuffer serialize(T input, ClassTag classTag)
{
Tuple2 tuple = (Tuple2) input;
PrestoSparkMutableRow row = tuple._2;
return row.getBuffer();
}
@Override
public T deserialize(ByteBuffer buffer, ClassTag classTag)
{
throw new UnsupportedOperationException("this method is never used by shuffle");
}
public T deserialize(byte[] array, int offset, int length, ClassTag classTag)
{
row.setArray(array);
row.setOffset(offset);
row.setLength(length);
return (T) tuple;
}
@Override
public T deserialize(ByteBuffer bytes, ClassLoader loader, ClassTag classTag)
{
throw new UnsupportedOperationException("this method is never used by shuffle");
}
}
public static class PrestoSparkShuffleSerializationStream
extends SerializationStream
{
private final DataOutputStream outputStream;
public PrestoSparkShuffleSerializationStream(OutputStream outputStream)
{
this.outputStream = new DataOutputStream(requireNonNull(outputStream, "outputStream is null"));
}
@Override
public SerializationStream writeKey(T key, ClassTag classTag)
{
// key is only needed to select partition
return this;
}
@Override
public SerializationStream writeValue(T value, ClassTag classTag)
{
PrestoSparkMutableRow row = (PrestoSparkMutableRow) value;
ByteBuffer buffer = row.getBuffer();
int length = buffer.remaining();
try {
outputStream.writeInt(length);
outputStream.write(buffer.array(), buffer.arrayOffset() + buffer.position(), length);
}
catch (IOException e) {
throw new UncheckedIOException(e);
}
return this;
}
@Override
public void flush()
{
try {
outputStream.flush();
}
catch (IOException e) {
throw new UncheckedIOException(e);
}
}
@Override
public void close()
{
try {
outputStream.close();
}
catch (IOException e) {
throw new UncheckedIOException(e);
}
}
@Override
public SerializationStream writeObject(T tuple, ClassTag classTag)
{
throw new UnsupportedOperationException("this method is never used by shuffle");
}
@Override
public SerializationStream writeAll(Iterator iterator, ClassTag classTag)
{
throw new UnsupportedOperationException("this method is never used by shuffle");
}
}
public static class PrestoSparkShuffleDeserializationStream
extends DeserializationStream
{
private final DataInputStream inputStream;
private final MutablePartitionId mutablePartitionId = new MutablePartitionId();
private final PrestoSparkMutableRow row = new PrestoSparkMutableRow();
private final Tuple2
© 2015 - 2025 Weber Informatics LLC | Privacy Policy