org.apache.beam.sdk.io.hadoop.WritableCoder Maven / Gradle / Ivy
/*
* 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.
*/
package org.apache.beam.sdk.io.hadoop;
import com.google.auto.service.AutoService;
import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Collections;
import java.util.List;
import java.util.Objects;
import org.apache.beam.sdk.coders.CannotProvideCoderException;
import org.apache.beam.sdk.coders.Coder;
import org.apache.beam.sdk.coders.CoderException;
import org.apache.beam.sdk.coders.CoderProvider;
import org.apache.beam.sdk.coders.CoderProviderRegistrar;
import org.apache.beam.sdk.coders.CustomCoder;
import org.apache.beam.sdk.coders.DefaultCoder;
import org.apache.beam.sdk.values.TypeDescriptor;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Writable;
/**
* A {@code WritableCoder} is a {@link Coder} for a Java class that implements {@link Writable}.
*
* To use, specify the coder type on a PCollection:
*
* {@code
* PCollection records =
* foo.apply(...).setCoder(WritableCoder.of(MyRecord.class));
* }
*
*
* @param the type of elements handled by this coder.
*/
public class WritableCoder extends CustomCoder {
private static final long serialVersionUID = 0L;
/**
* Returns a {@code WritableCoder} instance for the provided element class.
* @param the element type
*/
public static WritableCoder of(Class clazz) {
return new WritableCoder<>(clazz);
}
private final Class type;
public WritableCoder(Class type) {
this.type = type;
}
@Override
public void encode(T value, OutputStream outStream) throws IOException {
value.write(new DataOutputStream(outStream));
}
@SuppressWarnings("unchecked")
@Override
public T decode(InputStream inStream) throws IOException {
try {
if (type == NullWritable.class) {
// NullWritable has no default constructor
return (T) NullWritable.get();
}
T t = type.newInstance();
t.readFields(new DataInputStream(inStream));
return t;
} catch (InstantiationException | IllegalAccessException e) {
throw new CoderException("unable to deserialize record", e);
}
}
@Override
public List> getCoderArguments() {
return Collections.emptyList();
}
@Override
public void verifyDeterministic() throws NonDeterministicException {
throw new NonDeterministicException(this,
"Hadoop Writable may be non-deterministic.");
}
@Override
public boolean equals(Object other) {
if (other == this) {
return true;
}
if (!(other instanceof WritableCoder)) {
return false;
}
WritableCoder that = (WritableCoder) other;
return Objects.equals(this.type, that.type);
}
@Override
public int hashCode() {
return type.hashCode();
}
/**
* Returns a {@link CoderProvider} which uses the {@link WritableCoder} for Hadoop
* {@link Writable writable types}.
*
* This method is invoked reflectively from {@link DefaultCoder}.
*/
public static CoderProvider getCoderProvider() {
return new WritableCoderProvider();
}
/**
* A {@link CoderProviderRegistrar} which registers a {@link CoderProvider} which can handle
* {@link Writable writable types}.
*/
@AutoService(CoderProviderRegistrar.class)
public static class WritableCoderProviderRegistrar implements CoderProviderRegistrar {
@Override
public List getCoderProviders() {
return Collections.singletonList(getCoderProvider());
}
}
/**
* A {@link CoderProvider} for Hadoop {@link Writable writable types}.
*/
private static class WritableCoderProvider extends CoderProvider {
private static final TypeDescriptor WRITABLE_TYPE = new TypeDescriptor() {};
@Override
public Coder coderFor(TypeDescriptor typeDescriptor,
List> componentCoders) throws CannotProvideCoderException {
if (!typeDescriptor.isSubtypeOf(WRITABLE_TYPE)) {
throw new CannotProvideCoderException(
String.format(
"Cannot provide %s because %s does not implement the interface %s",
WritableCoder.class.getSimpleName(),
typeDescriptor,
Writable.class.getName()));
}
try {
@SuppressWarnings("unchecked")
Coder coder = WritableCoder.of((Class) typeDescriptor.getRawType());
return coder;
} catch (IllegalArgumentException e) {
throw new CannotProvideCoderException(e);
}
}
}
}