Search in sources :

Example 1 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class CoderRegistry method verifyCompatible.

/**
 * Returns {@code true} if the given {@link Coder} can possibly encode elements of the given type.
 */
@VisibleForTesting
static <T, CoderT extends Coder<T>, CandidateT> void verifyCompatible(CoderT coder, Type candidateType) throws IncompatibleCoderException {
    // Various representations of the coder's class
    @SuppressWarnings("unchecked") Class<CoderT> coderClass = (Class<CoderT>) coder.getClass();
    TypeDescriptor<CoderT> coderDescriptor = TypeDescriptor.of(coderClass);
    // Various representations of the actual coded type
    @SuppressWarnings("unchecked") TypeDescriptor<T> codedDescriptor = CoderUtils.getCodedType(coderDescriptor);
    @SuppressWarnings("unchecked") Class<T> codedClass = (Class<T>) codedDescriptor.getRawType();
    Type codedType = codedDescriptor.getType();
    // Various representations of the candidate type
    @SuppressWarnings("unchecked") TypeDescriptor<CandidateT> candidateDescriptor = (TypeDescriptor<CandidateT>) TypeDescriptor.of(candidateType);
    @SuppressWarnings("unchecked") Class<CandidateT> candidateClass = (Class<CandidateT>) candidateDescriptor.getRawType();
    // to erasure, then we cannot rule it out.
    if (candidateType instanceof TypeVariable) {
        return;
    }
    // coder compatibility
    if (!codedClass.isAssignableFrom(candidateClass)) {
        throw new IncompatibleCoderException(String.format("Cannot encode elements of type %s with coder %s because the" + " coded type %s is not assignable from %s", candidateType, coder, codedClass, candidateType), coder, candidateType);
    }
    // we have established that this is a covariant upcast... though
    // coders are invariant, we are just checking one direction
    @SuppressWarnings("unchecked") TypeDescriptor<T> candidateOkDescriptor = (TypeDescriptor<T>) candidateDescriptor;
    // compatible.
    if ((codedType instanceof ParameterizedType) && !isNullOrEmpty(coder.getCoderArguments())) {
        ParameterizedType parameterizedSupertype = (ParameterizedType) candidateOkDescriptor.getSupertype(codedClass).getType();
        Type[] typeArguments = parameterizedSupertype.getActualTypeArguments();
        List<? extends Coder<?>> typeArgumentCoders = coder.getCoderArguments();
        if (typeArguments.length < typeArgumentCoders.size()) {
            throw new IncompatibleCoderException(String.format("Cannot encode elements of type %s with coder %s:" + " the generic supertype %s has %s type parameters, which is less than the" + " number of coder arguments %s has (%s).", candidateOkDescriptor, coder, parameterizedSupertype, typeArguments.length, coder, typeArgumentCoders.size()), coder, candidateOkDescriptor.getType());
        }
        for (int i = 0; i < typeArgumentCoders.size(); i++) {
            try {
                Coder<?> typeArgumentCoder = typeArgumentCoders.get(i);
                verifyCompatible(typeArgumentCoder, candidateDescriptor.resolveType(typeArguments[i]).getType());
            } catch (IncompatibleCoderException exn) {
                throw new IncompatibleCoderException(String.format("Cannot encode elements of type %s with coder %s" + " because some component coder is incompatible", candidateType, coder), coder, candidateType, exn);
            }
        }
    }
}
Also used : ParameterizedType(java.lang.reflect.ParameterizedType) WildcardType(java.lang.reflect.WildcardType) ParameterizedType(java.lang.reflect.ParameterizedType) Type(java.lang.reflect.Type) TypeDescriptor(org.apache.beam.sdk.values.TypeDescriptor) TypeVariable(java.lang.reflect.TypeVariable) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Example 2 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class AvroSource method readMetadataFromFile.

/**
 * Reads the {@link AvroMetadata} from the header of an Avro file.
 *
 * <p>This method parses the header of an Avro <a
 * href="https://avro.apache.org/docs/1.7.7/spec.html#Object+Container+Files">Object Container
 * File</a>.
 *
 * @throws IOException if the file is an invalid format.
 */
@VisibleForTesting
static AvroMetadata readMetadataFromFile(ResourceId fileResource) throws IOException {
    String codec = null;
    String schemaString = null;
    byte[] syncMarker;
    try (InputStream stream = Channels.newInputStream(FileSystems.open(fileResource))) {
        BinaryDecoder decoder = DecoderFactory.get().binaryDecoder(stream, null);
        // The header of an object container file begins with a four-byte magic number, followed
        // by the file metadata (including the schema and codec), encoded as a map. Finally, the
        // header ends with the file's 16-byte sync marker.
        // See https://avro.apache.org/docs/1.7.7/spec.html#Object+Container+Files for details on
        // the encoding of container files.
        // Read the magic number.
        byte[] magic = new byte[DataFileConstants.MAGIC.length];
        decoder.readFixed(magic);
        if (!Arrays.equals(magic, DataFileConstants.MAGIC)) {
            throw new IOException("Missing Avro file signature: " + fileResource);
        }
        // Read the metadata to find the codec and schema.
        ByteBuffer valueBuffer = ByteBuffer.allocate(512);
        long numRecords = decoder.readMapStart();
        while (numRecords > 0) {
            for (long recordIndex = 0; recordIndex < numRecords; recordIndex++) {
                String key = decoder.readString();
                // readBytes() clears the buffer and returns a buffer where:
                // - position is the start of the bytes read
                // - limit is the end of the bytes read
                valueBuffer = decoder.readBytes(valueBuffer);
                byte[] bytes = new byte[valueBuffer.remaining()];
                valueBuffer.get(bytes);
                if (key.equals(DataFileConstants.CODEC)) {
                    codec = new String(bytes, StandardCharsets.UTF_8);
                } else if (key.equals(DataFileConstants.SCHEMA)) {
                    schemaString = new String(bytes, StandardCharsets.UTF_8);
                }
            }
            numRecords = decoder.mapNext();
        }
        if (codec == null) {
            codec = DataFileConstants.NULL_CODEC;
        }
        // Finally, read the sync marker.
        syncMarker = new byte[DataFileConstants.SYNC_SIZE];
        decoder.readFixed(syncMarker);
    }
    checkState(schemaString != null, "No schema present in Avro file metadata %s", fileResource);
    return new AvroMetadata(syncMarker, codec, schemaString);
}
Also used : PushbackInputStream(java.io.PushbackInputStream) ObjectInputStream(java.io.ObjectInputStream) ByteArrayInputStream(java.io.ByteArrayInputStream) XZCompressorInputStream(org.apache.commons.compress.compressors.xz.XZCompressorInputStream) CountingInputStream(org.apache.commons.compress.utils.CountingInputStream) InflaterInputStream(java.util.zip.InflaterInputStream) SnappyCompressorInputStream(org.apache.commons.compress.compressors.snappy.SnappyCompressorInputStream) BZip2CompressorInputStream(org.apache.commons.compress.compressors.bzip2.BZip2CompressorInputStream) InputStream(java.io.InputStream) IOException(java.io.IOException) ByteBuffer(java.nio.ByteBuffer) BinaryDecoder(org.apache.avro.io.BinaryDecoder) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Example 3 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class BeamRowToStorageApiProto method descriptorSchemaFromBeamSchema.

@VisibleForTesting
static DescriptorProto descriptorSchemaFromBeamSchema(Schema schema) {
    Preconditions.checkState(schema.getFieldCount() > 0);
    DescriptorProto.Builder descriptorBuilder = DescriptorProto.newBuilder();
    // Create a unique name for the descriptor ('-' characters cannot be used).
    descriptorBuilder.setName("D" + UUID.randomUUID().toString().replace("-", "_"));
    int i = 1;
    List<DescriptorProto> nestedTypes = Lists.newArrayList();
    for (Field field : schema.getFields()) {
        FieldDescriptorProto.Builder fieldDescriptorProtoBuilder = fieldDescriptorFromBeamField(field, i++, nestedTypes);
        descriptorBuilder.addField(fieldDescriptorProtoBuilder);
    }
    nestedTypes.forEach(descriptorBuilder::addNestedType);
    return descriptorBuilder.build();
}
Also used : Field(org.apache.beam.sdk.schemas.Schema.Field) FieldDescriptorProto(com.google.protobuf.DescriptorProtos.FieldDescriptorProto) FileDescriptorProto(com.google.protobuf.DescriptorProtos.FileDescriptorProto) DescriptorProto(com.google.protobuf.DescriptorProtos.DescriptorProto) FieldDescriptorProto(com.google.protobuf.DescriptorProtos.FieldDescriptorProto) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Example 4 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class GlobalWatermarkHolder method clear.

@VisibleForTesting
public static synchronized void clear() {
    sourceTimes.clear();
    lastWatermarkedBatchTime = 0;
    writeLocalWatermarkCopy(null);
    final SparkEnv sparkEnv = SparkEnv.get();
    if (sparkEnv != null) {
        final BlockManager blockManager = sparkEnv.blockManager();
        blockManager.removeBlock(WATERMARKS_BLOCK_ID, true);
    }
}
Also used : BlockManager(org.apache.spark.storage.BlockManager) SparkEnv(org.apache.spark.SparkEnv) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Example 5 with VisibleForTesting

use of org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting in project beam by apache.

the class BeamFnStatusClient method getCacheStats.

@VisibleForTesting
String getCacheStats() {
    StringJoiner cacheStats = new StringJoiner("\n");
    cacheStats.add("========== CACHE STATS ==========");
    cacheStats.add(cache.describeStats());
    return cacheStats.toString();
}
Also used : StringJoiner(java.util.StringJoiner) VisibleForTesting(org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)

Aggregations

VisibleForTesting (org.apache.beam.vendor.guava.v26_0_jre.com.google.common.annotations.VisibleForTesting)81 ArrayList (java.util.ArrayList)18 IOException (java.io.IOException)17 ParameterizedType (java.lang.reflect.ParameterizedType)15 Type (java.lang.reflect.Type)15 Parameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter)14 BundleFinalizerParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.BundleFinalizerParameter)14 PipelineOptionsParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.PipelineOptionsParameter)14 RestrictionParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.RestrictionParameter)14 RestrictionTrackerParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.RestrictionTrackerParameter)14 SchemaElementParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.SchemaElementParameter)14 StateParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.StateParameter)14 TimerFamilyParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimerFamilyParameter)14 TimerParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.TimerParameter)14 WatermarkEstimatorParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.WatermarkEstimatorParameter)14 WatermarkEstimatorStateParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.WatermarkEstimatorStateParameter)14 WindowParameter (org.apache.beam.sdk.transforms.reflect.DoFnSignature.Parameter.WindowParameter)14 TypeParameter (org.apache.beam.sdk.values.TypeParameter)14 DoFn (org.apache.beam.sdk.transforms.DoFn)10 Map (java.util.Map)7