Java Code Examples for org.apache.beam.sdk.values.Row#getFieldCount()

@SuppressWarnings("unchecked")
static void encodeDelegate(
    Coder[] coders, Row value, OutputStream outputStream, boolean hasNullableFields)
    throws IOException {
  checkState(value.getFieldCount() == value.getSchema().getFieldCount());

  // Encode the field count. This allows us to handle compatible schema changes.
  VAR_INT_CODER.encode(value.getFieldCount(), outputStream);
  // Encode a bitmap for the null fields to save having to encode a bunch of nulls.
  NULL_LIST_CODER.encode(scanNullFields(value, hasNullableFields), outputStream);
  for (int idx = 0; idx < value.getFieldCount(); ++idx) {
    Object fieldValue = value.getValue(idx);
    if (value.getValue(idx) != null) {
      coders[idx].encode(fieldValue, outputStream);
    }
  }
}
 

private static boolean rowsEquivalent(Row expected, Row actual) {
  if (!actual.getSchema().equivalent(expected.getSchema())) {
    return false;
  }
  if (expected.getFieldCount() != actual.getFieldCount()) {
    return false;
  }
  for (int i = 0; i < expected.getFieldCount(); ++i) {
    Field field = expected.getSchema().getField(i);
    int actualIndex = actual.getSchema().indexOf(field.getName());
    if (!fieldsEquivalent(
        expected.getValue(i), actual.getValue(actualIndex), field.getType())) {
      return false;
    }
  }
  return true;
}
 

public static String beamRow2CsvLine(Row row, CSVFormat csvFormat) {
  StringWriter writer = new StringWriter();
  try (CSVPrinter printer = csvFormat.print(writer)) {
    for (int i = 0; i < row.getFieldCount(); i++) {
      printer.print(row.getBaseValue(i, Object.class).toString());
    }
    printer.println();
  } catch (IOException e) {
    throw new IllegalArgumentException("encodeRecord failed!", e);
  }
  return writer.toString();
}
 

public static Value beamRowToZetaSqlStructValue(Row row, Schema schema) {
  List<Value> values = new ArrayList<>(row.getFieldCount());

  for (int i = 0; i < row.getFieldCount(); i++) {
    values.add(
        javaObjectToZetaSqlValue(
            row.getBaseValue(i, Object.class), schema.getField(i).getType()));
  }
  return Value.createStructValue(beamSchemaToZetaSqlStructType(schema), values);
}
 

private static BitSet scanNullFields(Row row, boolean hasNullableFields) {
  BitSet nullFields = new BitSet(row.getFieldCount());
  if (hasNullableFields) {
    for (int idx = 0; idx < row.getFieldCount(); ++idx) {
      if (row.getValue(idx) == null) {
        nullFields.set(idx);
      }
    }
  }
  return nullFields;
}
 

private static Row fillNewFields(Row row, AddFieldsInformation addFieldsInformation) {
  Schema outputSchema = checkNotNull(addFieldsInformation.getOutputFieldType().getRowSchema());

  List<Object> newValues = Lists.newArrayListWithCapacity(outputSchema.getFieldCount());
  for (int i = 0; i < row.getFieldCount(); ++i) {
    AddFieldsInformation nested = addFieldsInformation.getNestedNewValues().get(i);
    if (nested != null) {
      // New fields were added to nested subfields of this value. Recursively fill them out
      // before adding to the new row.
      Object newValue = fillNewFields(row.getValue(i), nested.getOutputFieldType(), nested);
      newValues.add(newValue);
    } else {
      // Nothing changed. Just copy the old value into the new row.
      newValues.add(row.getValue(i));
    }
  }
  // If there are brand new simple (i.e. have no nested values) fields at this level, then add
  // the default values for all of them.
  newValues.addAll(addFieldsInformation.getDefaultValues());
  // If we are creating new recursive fields, populate new values for them here.
  for (int i = newValues.size(); i < addFieldsInformation.getNestedNewValues().size(); ++i) {
    AddFieldsInformation newNestedField = addFieldsInformation.getNestedNewValues().get(i);
    if (newNestedField != null) {
      newValues.add(fillNewFields(null, newNestedField.getOutputFieldType(), newNestedField));
    }
  }

  return Row.withSchema(outputSchema).attachValues(newValues);
}
 

@SuppressWarnings("unchecked")
public <ValueT> ValueT fromRow(
    Row row, Class<ValueT> clazz, Factory<List<FieldValueTypeInformation>> typeFactory) {
  if (row instanceof RowWithGetters) {
    Object target = ((RowWithGetters) row).getGetterTarget();
    if (target.getClass().equals(clazz)) {
      // Efficient path: simply extract the underlying object instead of creating a new one.
      return (ValueT) target;
    }
  }

  Object[] params = new Object[row.getFieldCount()];
  Schema schema = row.getSchema();
  List<FieldValueTypeInformation> typeInformations = typeFactory.create(clazz, schema);
  checkState(
      typeInformations.size() == row.getFieldCount(),
      "Did not have a matching number of type informations and fields.");

  for (int i = 0; i < row.getFieldCount(); ++i) {
    FieldType type = schema.getField(i).getType();
    FieldValueTypeInformation typeInformation = checkNotNull(typeInformations.get(i));
    params[i] =
        fromValue(
            type, row.getValue(i), typeInformation.getRawType(), typeInformation, typeFactory);
  }

  SchemaUserTypeCreator creator = schemaTypeCreatorFactory.create(clazz, schema);
  return (ValueT) creator.create(params);
}
 

@Override
public Value toInputType(Row base) {
  EnumerationType.Value caseType = null;
  Object oneOfValue = null;
  for (int i = 0; i < base.getFieldCount(); ++i) {
    Object value = base.getValue(i);
    if (value != null) {
      checkArgument(caseType == null, "More than one field set in union " + this);
      caseType = enumerationType.valueOf(oneOfSchema.getField(i).getName());
      oneOfValue = value;
    }
  }
  checkNotNull(oneOfValue, "No value set in union" + this);
  return createValue(caseType, oneOfValue);
}
 

public static SchemaApi.Row rowToProto(Row row) {
  SchemaApi.Row.Builder builder = SchemaApi.Row.newBuilder();
  for (int i = 0; i < row.getFieldCount(); ++i) {
    builder.addValues(fieldValueToProto(row.getSchema().getField(i).getType(), row.getValue(i)));
  }
  return builder.build();
}
 

/** Convert a BigQuery TableRow to a Beam Row. */
public static TableRow toTableRow(Row row) {
  TableRow output = new TableRow();
  for (int i = 0; i < row.getFieldCount(); i++) {
    Object value = row.getValue(i);
    Field schemaField = row.getSchema().getField(i);
    output = output.set(schemaField.getName(), fromBeamField(schemaField.getType(), value));
  }
  return output;
}