Java Code Examples for org.apache.beam.sdk.schemas.Schema#FieldType

public static boolean deepEquals(Object a, Object b, Schema.FieldType fieldType) {
  if (a == null || b == null) {
    return a == b;
  } else if (fieldType.getTypeName() == TypeName.LOGICAL_TYPE) {
    return deepEquals(a, b, fieldType.getLogicalType().getBaseType());
  } else if (fieldType.getTypeName() == Schema.TypeName.BYTES) {
    return Arrays.equals((byte[]) a, (byte[]) b);
  } else if (fieldType.getTypeName() == TypeName.ARRAY) {
    return deepEqualsForCollection(
        (Collection<Object>) a, (Collection<Object>) b, fieldType.getCollectionElementType());
  } else if (fieldType.getTypeName() == TypeName.ITERABLE) {
    return deepEqualsForIterable(
        (Iterable<Object>) a, (Iterable<Object>) b, fieldType.getCollectionElementType());
  } else if (fieldType.getTypeName() == Schema.TypeName.MAP) {
    return deepEqualsForMap(
        (Map<Object, Object>) a, (Map<Object, Object>) b, fieldType.getMapValueType());
  } else {
    return Objects.equals(a, b);
  }
}
 

@Override
public Integer accept(Context context, Schema.FieldType left, Schema.FieldType right) {

  if (left.getTypeName() != right.getTypeName()) {
    return 0;
  }

  if (left.getTypeName() == Schema.TypeName.ROW) {
    return 0;
  }

  if (left.getTypeName() == Schema.TypeName.ARRAY) {
    return 0;
  }

  if (left.getTypeName() == Schema.TypeName.MAP) {
    return 0;
  }

  return 1;
}
 

static boolean deepEqualsForIterable(
    Iterable<Object> a, Iterable<Object> b, Schema.FieldType elementType) {
  if (a == b) {
    return true;
  }
  Iterator<Object> bIter = b.iterator();
  for (Object currentA : a) {
    if (!bIter.hasNext()) {
      return false;
    }
    if (!deepEquals(currentA, bIter.next(), elementType)) {
      return false;
    }
  }
  return !bIter.hasNext();
}
 

static boolean iterablesEquivalent(
    Iterable<Object> expected, Iterable<Object> actual, Schema.FieldType elementType) {
  if (expected == actual) {
    return true;
  }
  Iterator<Object> actualIter = actual.iterator();
  for (Object currentExpected : expected) {
    if (!actualIter.hasNext()) {
      return false;
    }
    if (!fieldsEquivalent(currentExpected, actualIter.next(), elementType)) {
      return false;
    }
  }
  return !actualIter.hasNext();
}
 

/** {@link CombineFn} for Sum based on {@link Sum} and {@link Combine.BinaryCombineFn}. */
static CombineFn createSum(Schema.FieldType fieldType) {
  switch (fieldType.getTypeName()) {
    case INT32:
      return Sum.ofIntegers();
    case INT16:
      return new ShortSum();
    case BYTE:
      return new ByteSum();
    case INT64:
      return Sum.ofLongs();
    case FLOAT:
      return new FloatSum();
    case DOUBLE:
      return Sum.ofDoubles();
    case DECIMAL:
      return new BigDecimalSum();
    default:
      throw new UnsupportedOperationException(
          String.format("[%s] is not supported in SUM", fieldType));
  }
}
 

private static Object convertMapStrict(
    Map<CharSequence, Object> values,
    org.apache.avro.Schema valueAvroSchema,
    Schema.FieldType fieldType) {
  checkTypeName(fieldType.getTypeName(), Schema.TypeName.MAP, "map");
  checkNotNull(fieldType.getMapKeyType());
  checkNotNull(fieldType.getMapValueType());

  if (!fieldType.getMapKeyType().equals(Schema.FieldType.STRING)) {
    throw new IllegalArgumentException(
        "Can't convert 'string' map keys to " + fieldType.getMapKeyType());
  }

  Map<Object, Object> ret = new HashMap<>();

  for (Map.Entry<CharSequence, Object> value : values.entrySet()) {
    ret.put(
        convertStringStrict(value.getKey(), fieldType.getMapKeyType()),
        convertAvroFieldStrict(value.getValue(), valueAvroSchema, fieldType.getMapValueType()));
  }

  return ret;
}
 

/**
 * Infer a schema from a Java class.
 *
 * <p>Takes in a function to extract a list of field types from a class. Different callers may
 * have different strategies for extracting this list: e.g. introspecting public member variables,
 * public getter methods, or special annotations on the class.
 */
public static Schema schemaFromClass(
    Class<?> clazz, FieldValueTypeSupplier fieldValueTypeSupplier) {
  Schema.Builder builder = Schema.builder();
  for (FieldValueTypeInformation type : fieldValueTypeSupplier.get(clazz)) {
    Schema.FieldType fieldType = fieldFromType(type.getType(), fieldValueTypeSupplier);
    if (type.isNullable()) {
      builder.addNullableField(type.getName(), fieldType);
    } else {
      builder.addField(type.getName(), fieldType);
    }
  }
  return builder.build();
}
 

static boolean deepEqualsForCollection(
    Collection<Object> a, Collection<Object> b, Schema.FieldType elementType) {
  if (a == b) {
    return true;
  }

  if (a.size() != b.size()) {
    return false;
  }

  return deepEqualsForIterable(a, b, elementType);
}
 

static int deepHashCodeForIterable(Iterable<Object> a, Schema.FieldType elementType) {
  int h = 1;
  for (Object o : a) {
    h = 31 * h + deepHashCode(o, elementType);
  }

  return h;
}
 

@Override
public Schema.FieldType getArgumentType() {
  return Schema.FieldType.STRING;
}
 

@Override
public Schema.FieldType getBaseType() {
  return Schema.FieldType.INT64;
}
 

/** Converts AVRO schema to Beam field. */
private static Schema.FieldType toFieldType(TypeWithNullability type) {
  Schema.FieldType fieldType = null;
  org.apache.avro.Schema avroSchema = type.type;

  LogicalType logicalType = LogicalTypes.fromSchema(avroSchema);
  if (logicalType != null) {
    if (logicalType instanceof LogicalTypes.Decimal) {
      fieldType = FieldType.DECIMAL;
    } else if (logicalType instanceof LogicalTypes.TimestampMillis) {
      // TODO: There is a desire to move Beam schema DATETIME to a micros representation. When
      // this is done, this logical type needs to be changed.
      fieldType = FieldType.DATETIME;
    } else if (logicalType instanceof LogicalTypes.Date) {
      fieldType = FieldType.DATETIME;
    }
  }

  if (fieldType == null) {
    switch (type.type.getType()) {
      case RECORD:
        fieldType = Schema.FieldType.row(toBeamSchema(avroSchema));
        break;

      case ENUM:
        fieldType = FieldType.logicalType(EnumerationType.create(type.type.getEnumSymbols()));
        break;

      case ARRAY:
        Schema.FieldType elementType =
            toFieldType(new TypeWithNullability(avroSchema.getElementType()));
        fieldType = Schema.FieldType.array(elementType);
        break;

      case MAP:
        fieldType =
            Schema.FieldType.map(
                Schema.FieldType.STRING,
                toFieldType(new TypeWithNullability(avroSchema.getValueType())));
        break;

      case FIXED:
        fieldType = FixedBytesField.fromAvroType(type.type).toBeamType();
        break;

      case STRING:
        fieldType = Schema.FieldType.STRING;
        break;

      case BYTES:
        fieldType = Schema.FieldType.BYTES;
        break;

      case INT:
        fieldType = Schema.FieldType.INT32;
        break;

      case LONG:
        fieldType = Schema.FieldType.INT64;
        break;

      case FLOAT:
        fieldType = Schema.FieldType.FLOAT;
        break;

      case DOUBLE:
        fieldType = Schema.FieldType.DOUBLE;
        break;

      case BOOLEAN:
        fieldType = Schema.FieldType.BOOLEAN;
        break;

      case UNION:
        throw new IllegalArgumentException("Union types not yet supported");

      case NULL:
        throw new IllegalArgumentException("Can't convert 'null' to FieldType");

      default:
        throw new AssertionError("Unexpected AVRO Schema.Type: " + avroSchema.getType());
    }
  }
  fieldType = fieldType.withNullable(type.nullable);
  return fieldType;
}
 

private static org.apache.avro.Schema getFieldSchema(
    Schema.FieldType fieldType, String fieldName, String namespace) {
  org.apache.avro.Schema baseType;
  switch (fieldType.getTypeName()) {
    case BYTE:
    case INT16:
    case INT32:
      baseType = org.apache.avro.Schema.create(Type.INT);
      break;

    case INT64:
      baseType = org.apache.avro.Schema.create(Type.LONG);
      break;

    case DECIMAL:
      baseType =
          LogicalTypes.decimal(Integer.MAX_VALUE)
              .addToSchema(org.apache.avro.Schema.create(Type.BYTES));
      break;

    case FLOAT:
      baseType = org.apache.avro.Schema.create(Type.FLOAT);
      break;

    case DOUBLE:
      baseType = org.apache.avro.Schema.create(Type.DOUBLE);
      break;

    case STRING:
      baseType = org.apache.avro.Schema.create(Type.STRING);
      break;

    case DATETIME:
      // TODO: There is a desire to move Beam schema DATETIME to a micros representation. When
      // this is done, this logical type needs to be changed.
      baseType =
          LogicalTypes.timestampMillis().addToSchema(org.apache.avro.Schema.create(Type.LONG));
      break;

    case BOOLEAN:
      baseType = org.apache.avro.Schema.create(Type.BOOLEAN);
      break;

    case BYTES:
      baseType = org.apache.avro.Schema.create(Type.BYTES);
      break;

    case LOGICAL_TYPE:
      switch (fieldType.getLogicalType().getIdentifier()) {
        case FixedBytes.IDENTIFIER:
          FixedBytesField fixedBytesField =
              checkNotNull(FixedBytesField.fromBeamFieldType(fieldType));
          baseType = fixedBytesField.toAvroType("fixed", namespace + "." + fieldName);
          break;
        case EnumerationType.IDENTIFIER:
          EnumerationType enumerationType = fieldType.getLogicalType(EnumerationType.class);
          baseType =
              org.apache.avro.Schema.createEnum(fieldName, "", "", enumerationType.getValues());
          break;
        default:
          throw new RuntimeException(
              "Unhandled logical type " + fieldType.getLogicalType().getIdentifier());
      }
      break;

    case ARRAY:
    case ITERABLE:
      baseType =
          org.apache.avro.Schema.createArray(
              getFieldSchema(fieldType.getCollectionElementType(), fieldName, namespace));
      break;

    case MAP:
      if (fieldType.getMapKeyType().getTypeName().isStringType()) {
        // Avro only supports string keys in maps.
        baseType =
            org.apache.avro.Schema.createMap(
                getFieldSchema(fieldType.getMapValueType(), fieldName, namespace));
      } else {
        throw new IllegalArgumentException("Avro only supports maps with string keys");
      }
      break;

    case ROW:
      baseType = toAvroSchema(fieldType.getRowSchema(), fieldName, namespace);
      break;

    default:
      throw new IllegalArgumentException("Unexpected type " + fieldType);
  }
  return fieldType.getNullable() ? ReflectData.makeNullable(baseType) : baseType;
}
 

@VisibleForTesting
static Schema.FieldType numeric(int precision, int scale) {
  return Schema.FieldType.logicalType(
      FixedPrecisionNumeric.of(JDBCType.NUMERIC.getName(), precision, scale));
}
 

@Override
public Schema.FieldType getBaseType() {
  return Schema.FieldType.INT64;
}
 

private static Object convertRecordStrict(GenericRecord record, Schema.FieldType fieldType) {
  checkTypeName(fieldType.getTypeName(), Schema.TypeName.ROW, "record");
  return toBeamRowStrict(record, fieldType.getRowSchema());
}
 

private static Object convertStringStrict(CharSequence value, Schema.FieldType fieldType) {
  checkTypeName(fieldType.getTypeName(), Schema.TypeName.STRING, "string");
  return value.toString();
}
 

private static Object fillNewFields(
    Object original, Schema.FieldType fieldType, AddFieldsInformation addFieldsInformation) {
  switch (fieldType.getTypeName()) {
    case ROW:
      if (original == null) {
        original = Row.withSchema(fieldType.getRowSchema()).build();
      }
      return fillNewFields((Row) original, addFieldsInformation);

    case ARRAY:
    case ITERABLE:
      if (original == null) {
        return Collections.emptyList();
      }
      Iterable<Object> iterable = (Iterable<Object>) original;
      List<Object> filledList = new ArrayList<>(Iterables.size(iterable));
      Schema.FieldType elementType = fieldType.getCollectionElementType();
      AddFieldsInformation elementAddFieldInformation =
          addFieldsInformation.toBuilder().setOutputFieldType(elementType).build();
      for (Object element : iterable) {
        filledList.add(fillNewFields(element, elementType, elementAddFieldInformation));
      }
      return filledList;

    case MAP:
      if (original == null) {
        return Collections.emptyMap();
      }
      Map<Object, Object> originalMap = (Map<Object, Object>) original;
      Map<Object, Object> filledMap = Maps.newHashMapWithExpectedSize(originalMap.size());
      Schema.FieldType mapValueType = fieldType.getMapValueType();
      AddFieldsInformation mapValueAddFieldInformation =
          addFieldsInformation.toBuilder().setOutputFieldType(mapValueType).build();
      for (Map.Entry<Object, Object> entry : originalMap.entrySet()) {
        filledMap.put(
            entry.getKey(),
            fillNewFields(entry.getValue(), mapValueType, mapValueAddFieldInformation));
      }
      return filledMap;

    default:
      throw new RuntimeException("Unexpected field type");
  }
}
 

private static Object convertLongStrict(Long value, Schema.FieldType fieldType) {
  checkTypeName(fieldType.getTypeName(), Schema.TypeName.INT64, "long");
  return value;
}
 

MapConvert(ProtoDynamicMessageSchema protoSchema, Schema.Field field) {
  super(field);
  Schema.FieldType fieldType = field.getType();
  key = protoSchema.createConverter(Schema.Field.of("KEY", fieldType.getMapKeyType()));
  value = protoSchema.createConverter(Schema.Field.of("VALUE", fieldType.getMapValueType()));
}