org.apache.flink.table.types.logical.RowType Java Examples

private SerializationRuntimeConverter createRowConverter(RowType type) {
	LogicalType[] fieldTypes = type.getFields().stream()
		.map(RowType.RowField::getType)
		.toArray(LogicalType[]::new);
	final String[] fieldNames = type.getFieldNames().toArray(new String[0]);
	final RowFieldConverter[] fieldConverters = Arrays.stream(fieldTypes)
		.map(this::createNullableRowFieldConverter)
		.toArray(RowFieldConverter[]::new);
	final int rowArity = type.getFieldCount();
	return (csvMapper, container, row) -> {
		// top level reuses the object node container
		final ObjectNode objectNode = (ObjectNode) container;
		for (int i = 0; i < rowArity; i++) {
			objectNode.set(
				fieldNames[i],
				fieldConverters[i].convert(csvMapper, container, row, i));
		}
		return objectNode;
	};
}
 

@Override
public FlinkFnApi.Schema.FieldType visit(RowType rowType) {
	FlinkFnApi.Schema.FieldType.Builder builder =
		FlinkFnApi.Schema.FieldType.newBuilder()
			.setTypeName(FlinkFnApi.Schema.TypeName.ROW)
			.setNullable(rowType.isNullable());

	FlinkFnApi.Schema.Builder schemaBuilder = FlinkFnApi.Schema.newBuilder();
	for (RowType.RowField field : rowType.getFields()) {
		schemaBuilder.addFields(
			FlinkFnApi.Schema.Field.newBuilder()
				.setName(field.getName())
				.setDescription(field.getDescription().orElse(EMPTY_STRING))
				.setType(field.getType().accept(this))
				.build());
	}
	builder.setRowSchema(schemaBuilder.build());
	return builder.build();
}
 

@Override
public AbstractPythonTableFunctionRunner<Row> createPythonTableFunctionRunner(
	PythonFunctionInfo pythonFunctionInfo,
	RowType inputType,
	RowType outputType) throws Exception {
	final FnDataReceiver<byte[]> dummyReceiver = input -> {
		// ignore the execution results
	};

	final PythonEnvironmentManager environmentManager = createTestEnvironmentManager();

	return new PythonTableFunctionRunner(
		"testPythonRunner",
		dummyReceiver,
		pythonFunctionInfo,
		environmentManager,
		inputType,
		outputType,
		Collections.emptyMap(),
		PythonTestUtils.createMockFlinkMetricContainer());
}
 

@Test
public void testHashSetMode() throws IOException {
	final int numMemSegments = needNumMemSegments(
			NUM_ENTRIES,
			rowLength(RowType.of(valueTypes)),
			rowLength(RowType.of(keyTypes)),
			PAGE_SIZE);
	int memorySize = numMemSegments * PAGE_SIZE;
	MemoryManager memoryManager = new MemoryManager(numMemSegments * PAGE_SIZE, 32);

	BytesHashMap table = new BytesHashMap(this, memoryManager,
			memorySize, keyTypes, new LogicalType[]{});
	Assert.assertTrue(table.isHashSetMode());

	final Random rnd = new Random(RANDOM_SEED);
	BinaryRow[] keys = getRandomizedInput(NUM_ENTRIES, rnd, true);
	verifyKeyInsert(keys, table);
	verifyKeyPresent(keys, table);
	table.free();
}
 

private static void writeCsvToStream(
		DataType[] types,
		RowData rowData,
		OutputStream stream) throws IOException {
	LogicalType[] fieldTypes = Arrays.stream(types)
			.map(DataType::getLogicalType)
			.toArray(LogicalType[]::new);
	DataFormatConverters.DataFormatConverter converter = DataFormatConverters.getConverterForDataType(
			TypeConversions.fromLogicalToDataType(RowType.of(fieldTypes)));

	Row row = (Row) converter.toExternal(rowData);
	StringBuilder builder = new StringBuilder();
	Object o;
	for (int i = 0; i < row.getArity(); i++) {
		if (i > 0) {
			builder.append(DEFAULT_FIELD_DELIMITER);
		}
		if ((o = row.getField(i)) != null) {
			builder.append(o);
		}
	}
	String str = builder.toString();
	stream.write(str.getBytes(StandardCharsets.UTF_8));
	stream.write(DEFAULT_LINE_DELIMITER.getBytes(StandardCharsets.UTF_8));
}
 

@Test
public void testBuildAndUpdate() throws Exception {
	final Random rnd = new Random(RANDOM_SEED);
	final BinaryRow[] rows = getRandomizedInput(NUM_ENTRIES, rnd, true);
	final int numMemSegments = needNumMemSegments(
			NUM_ENTRIES,
			rowLength(RowType.of(valueTypes)),
			rowLength(RowType.of(keyTypes)),
			PAGE_SIZE);
	int memorySize = numMemSegments * PAGE_SIZE;

	MemoryManager memoryManager = new MemoryManager(memorySize, 32);

	BytesHashMap table = new BytesHashMap(this, memoryManager,
			memorySize, keyTypes, valueTypes);

	List<BinaryRow> expected = new ArrayList<>(NUM_ENTRIES);
	verifyInsertAndUpdate(rnd, rows, expected, table);
	verifyRetrieve(table, rows, expected);
	table.free();
}
 

@Test
public void testJdbcFieldsProjection() {
	Map<String, String> properties = getBasicProperties();
	properties.put("connector.driver", "org.apache.derby.jdbc.EmbeddedDriver");
	properties.put("connector.username", "user");
	properties.put("connector.password", "pass");

	final TableSource<?> actual = ((JdbcTableSource) TableFactoryService
		.find(StreamTableSourceFactory.class, properties)
		.createStreamTableSource(properties))
		.projectFields(new int[] {0, 2});

	List<DataType> projectedFields = actual.getProducedDataType().getChildren();
	assertEquals(Arrays.asList(DataTypes.INT(), DataTypes.DOUBLE()), projectedFields);

	// test jdbc table source description
	List<String> fieldNames = ((RowType) actual.getProducedDataType().getLogicalType()).getFieldNames();
	String expectedSourceDescription = actual.getClass().getSimpleName()
		+ "(" + String.join(", ", fieldNames.stream().toArray(String[]::new)) + ")";
	assertEquals(expectedSourceDescription, actual.explainSource());
}
 

@Test
public void testHashSetMode() throws IOException {
	final int numMemSegments = needNumMemSegments(
			NUM_ENTRIES,
			rowLength(RowType.of(valueTypes)),
			rowLength(RowType.of(keyTypes)),
			PAGE_SIZE);
	int memorySize = numMemSegments * PAGE_SIZE;
	MemoryManager memoryManager = MemoryManagerBuilder
		.newBuilder()
		.setMemorySize(numMemSegments * PAGE_SIZE)
		.build();

	BytesHashMap table = new BytesHashMap(this, memoryManager,
			memorySize, keyTypes, new LogicalType[]{});
	Assert.assertTrue(table.isHashSetMode());

	final Random rnd = new Random(RANDOM_SEED);
	BinaryRowData[] keys = getRandomizedInput(NUM_ENTRIES, rnd, true);
	verifyKeyInsert(keys, table);
	verifyKeyPresent(keys, table);
	table.free();
}
 

private Optional<BulkWriter.Factory<RowData>> createBulkWriterFactory(String[] partitionColumns,
		StorageDescriptor sd) {
	String serLib = sd.getSerdeInfo().getSerializationLib().toLowerCase();
	int formatFieldCount = tableSchema.getFieldCount() - partitionColumns.length;
	String[] formatNames = new String[formatFieldCount];
	LogicalType[] formatTypes = new LogicalType[formatFieldCount];
	for (int i = 0; i < formatFieldCount; i++) {
		formatNames[i] = tableSchema.getFieldName(i).get();
		formatTypes[i] = tableSchema.getFieldDataType(i).get().getLogicalType();
	}
	RowType formatType = RowType.of(formatTypes, formatNames);
	Configuration formatConf = new Configuration(jobConf);
	sd.getSerdeInfo().getParameters().forEach(formatConf::set);
	if (serLib.contains("parquet")) {
		return Optional.of(ParquetRowDataBuilder.createWriterFactory(
				formatType, formatConf, hiveVersion.startsWith("3.")));
	} else if (serLib.contains("orc")) {
		TypeDescription typeDescription = OrcSplitReaderUtil.logicalTypeToOrcType(formatType);
		return Optional.of(hiveShim.createOrcBulkWriterFactory(
				formatConf, typeDescription.toString(), formatTypes));
	} else {
		return Optional.empty();
	}
}
 

private static TableSchema expandCompositeType(FieldsDataType dataType) {
	DataType[] fieldDataTypes = dataType.getChildren().toArray(new DataType[0]);
	return dataType.getLogicalType().accept(new LogicalTypeDefaultVisitor<TableSchema>() {
		@Override
		public TableSchema visit(RowType rowType) {
			return expandCompositeType(rowType, fieldDataTypes);
		}

		@Override
		public TableSchema visit(StructuredType structuredType) {
			return expandCompositeType(structuredType, fieldDataTypes);
		}

		@Override
		public TableSchema visit(DistinctType distinctType) {
			return distinctType.getSourceType().accept(this);
		}

		@Override
		protected TableSchema defaultMethod(LogicalType logicalType) {
			throw new IllegalArgumentException("Expected a composite type");
		}
	});
}
 

private List<RowType.RowField> parseRowFields(TokenType endToken) {
	List<RowType.RowField> fields = new ArrayList<>();
	boolean isFirst = true;
	while (!hasNextToken(endToken)) {
		if (isFirst) {
			isFirst = false;
		} else {
			nextToken(TokenType.LIST_SEPARATOR);
		}
		nextToken(TokenType.IDENTIFIER);
		final String name = tokenAsString();
		final LogicalType type = parseTypeWithNullability();
		if (hasNextToken(TokenType.LITERAL_STRING)) {
			nextToken(TokenType.LITERAL_STRING);
			final String description = tokenAsString();
			fields.add(new RowType.RowField(name, type, description));
		} else {
			fields.add(new RowType.RowField(name, type));
		}
	}
	return fields;
}
 

@Override
public AbstractPythonTableFunctionRunner<RowData> createPythonTableFunctionRunner(
	PythonFunctionInfo pythonFunctionInfo,
	RowType inputType,
	RowType outputType) throws Exception {
	final FnDataReceiver<byte[]> dummyReceiver = input -> {
		// ignore the execution results
	};

	final PythonEnvironmentManager environmentManager = createTestEnvironmentManager();

	return new RowDataPythonTableFunctionRunner(
		"testPythonRunner",
		dummyReceiver,
		pythonFunctionInfo,
		environmentManager,
		inputType,
		outputType,
		Collections.emptyMap(),
		PythonTestUtils.createMockFlinkMetricContainer());
}
 

private static Field toArrowField(String fieldName, LogicalType logicalType) {
	FieldType fieldType = new FieldType(
		logicalType.isNullable(),
		logicalType.accept(LogicalTypeToArrowTypeConverter.INSTANCE),
		null);
	List<Field> children = null;
	if (logicalType instanceof ArrayType) {
		children = Collections.singletonList(toArrowField(
			"element", ((ArrayType) logicalType).getElementType()));
	} else if (logicalType instanceof RowType) {
		RowType rowType = (RowType) logicalType;
		children = new ArrayList<>(rowType.getFieldCount());
		for (RowType.RowField field : rowType.getFields()) {
			children.add(toArrowField(field.getName(), field.getType()));
		}
	}
	return new Field(fieldName, fieldType, children);
}
 

@Override
public AbstractGeneralPythonScalarFunctionRunner<Row> createPythonScalarFunctionRunner(
	final PythonFunctionInfo[] pythonFunctionInfos,
	RowType inputType,
	RowType outputType) {
	final FnDataReceiver<byte[]> dummyReceiver = input -> {
		// ignore the execution results
	};

	final PythonEnvironmentManager environmentManager = createTestEnvironmentManager();

	return new PythonScalarFunctionRunner(
		"testPythonRunner",
		dummyReceiver,
		pythonFunctionInfos,
		environmentManager,
		inputType,
		outputType,
		Collections.emptyMap(),
		PythonTestUtils.createMockFlinkMetricContainer());
}
 

public static FieldsDataType pulsarSourceSchema(SchemaInfo si) throws IncompatibleSchemaException {
    List<DataTypes.Field> mainSchema = new ArrayList<>();
    DataType dataType = si2SqlType(si);
    if (dataType instanceof FieldsDataType) {
        FieldsDataType fieldsDataType = (FieldsDataType) dataType;
        RowType rowType = (RowType) fieldsDataType.getLogicalType();
        rowType.getFieldNames().stream()
                .map(fieldName -> DataTypes.FIELD(fieldName, fieldsDataType.getFieldDataTypes().get(fieldName)))
                .forEach(mainSchema::add);
    } else {
        mainSchema.add(DataTypes.FIELD("value", dataType));
    }

    mainSchema.addAll(METADATA_FIELDS);
    return (FieldsDataType) DataTypes.ROW(mainSchema.toArray(new DataTypes.Field[0]));
}
 

@Test
public void testInvalidTimestampTypeAvroSchemaConversion() {
	RowType rowType = (RowType) TableSchema.builder()
		.field("a", DataTypes.STRING())
		.field("b", DataTypes.TIMESTAMP(9))
		.build().toRowDataType().getLogicalType();
	thrown.expect(IllegalArgumentException.class);
	thrown.expectMessage("Avro does not support TIMESTAMP type with precision: 9, " +
		"it only supports precision less than 3.");
	AvroSchemaConverter.convertToSchema(rowType);
}
 

private AbstractGeneralPythonScalarFunctionRunner<Row> createUDFRunner(
	JobBundleFactory jobBundleFactory, FnDataReceiver<byte[]> receiver) {
	PythonFunctionInfo[] pythonFunctionInfos = new PythonFunctionInfo[] {
		new PythonFunctionInfo(
			DummyPythonFunction.INSTANCE,
			new Integer[]{0})
	};

	RowType rowType = new RowType(Collections.singletonList(new RowType.RowField("f1", new BigIntType())));

	final PythonEnvironmentManager environmentManager = createTestEnvironmentManager();

	return new PassThroughPythonScalarFunctionRunner<Row>(
		"testPythonRunner",
		receiver,
		pythonFunctionInfos,
		environmentManager,
		rowType,
		rowType,
		Collections.emptyMap(),
		jobBundleFactory,
		PythonTestUtils.createMockFlinkMetricContainer()) {
		@Override
		public TypeSerializer<Row> getInputTypeSerializer() {
			return (RowSerializer) PythonTypeUtils.toFlinkTypeSerializer(getInputType());
		}
	};
}
 

public AbstractPythonScalarFunctionOperator<CRow, CRow, Row> getTestOperator(
	Configuration config,
	PythonFunctionInfo[] scalarFunctions,
	RowType inputType,
	RowType outputType,
	int[] udfInputOffsets,
	int[] forwardedFields) {
	return new PassThroughArrowPythonScalarFunctionOperator(
		config, scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFields);
}
 

public RowDataArrowPythonScalarFunctionOperator(
	Configuration config,
	PythonFunctionInfo[] scalarFunctions,
	RowType inputType,
	RowType outputType,
	int[] udfInputOffsets,
	int[] forwardedFields) {
	super(config, scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFields);
}
 

public PythonScalarFunctionFlatMap(
	Configuration config,
	PythonFunctionInfo[] scalarFunctions,
	RowType inputType,
	RowType outputType,
	int[] udfInputOffsets,
	int[] forwardedFields) {
	super(config, scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFields);
}
 

@Override
public Optional<BulkWriter.Factory<RowData>> createBulkWriterFactory(WriterContext context) {
	return Optional.of(ParquetRowDataBuilder.createWriterFactory(
			RowType.of(Arrays.stream(context.getFormatFieldTypes())
							.map(DataType::getLogicalType)
							.toArray(LogicalType[]::new),
					context.getFormatFieldNames()),
			getParquetConfiguration(context.getFormatOptions()),
			context.getFormatOptions().get(UTC_TIMEZONE)));
}
 

private static JdbcBatchStatementExecutor<RowData> createKeyedRowExecutor(JdbcDialect dialect, int[] pkFields, LogicalType[] pkTypes, String sql, LogicalType[] logicalTypes) {
	final JdbcRowConverter rowConverter = dialect.getRowConverter(RowType.of(logicalTypes));
	final Function<RowData, RowData>  keyExtractor = createRowKeyExtractor(logicalTypes, pkFields);
	return JdbcBatchStatementExecutor.keyed(
		sql,
		keyExtractor,
		(st, record) -> rowConverter
			.toExternal(keyExtractor.apply(record), st));
}
 

public ParquetRowDataWriter(
		RecordConsumer recordConsumer,
		RowType rowType,
		GroupType schema,
		boolean utcTimestamp) {
	this.recordConsumer = recordConsumer;
	this.utcTimestamp = utcTimestamp;

	this.filedWriters = new FieldWriter[rowType.getFieldCount()];
	this.fieldNames = rowType.getFieldNames().toArray(new String[0]);
	for (int i = 0; i < rowType.getFieldCount(); i++) {
		this.filedWriters[i] = createWriter(rowType.getTypeAt(i), schema.getType(i));
	}
}
 

@Override
public void open(Configuration parameters) throws Exception {
	super.open(parameters);

	RowTypeInfo forwardedInputTypeInfo = new RowTypeInfo(
		Arrays.stream(forwardedFields)
			.mapToObj(i -> inputType.getFields().get(i))
			.map(RowType.RowField::getType)
			.map(TypeConversions::fromLogicalToDataType)
			.map(TypeConversions::fromDataTypeToLegacyInfo)
			.toArray(TypeInformation[]::new));
	forwardedInputSerializer = forwardedInputTypeInfo.createSerializer(getRuntimeContext().getExecutionConfig());
}
 

private Object value3(LogicalType type, Random rnd) {
	switch (type.getTypeRoot()) {
		case BOOLEAN:
			return true;
		case TINYINT:
			return Byte.MAX_VALUE;
		case SMALLINT:
			return Short.MAX_VALUE;
		case INTEGER:
			return Integer.MAX_VALUE;
		case BIGINT:
			return Long.MAX_VALUE;
		case FLOAT:
			return Float.MAX_VALUE;
		case DOUBLE:
			return Double.MAX_VALUE;
		case VARCHAR:
			return BinaryString.fromString(RandomStringUtils.random(100));
		case DECIMAL:
			DecimalType decimalType = (DecimalType) type;
			return Decimal.fromBigDecimal(new BigDecimal(Integer.MAX_VALUE),
					decimalType.getPrecision(), decimalType.getScale());
		case ARRAY:
		case VARBINARY:
			byte[] bytes = new byte[rnd.nextInt(100) + 100];
			rnd.nextBytes(bytes);
			return type instanceof VarBinaryType ? bytes : BinaryArray.fromPrimitiveArray(bytes);
		case ROW:
			RowType rowType = (RowType) type;
			if (rowType.getFields().get(0).getType().getTypeRoot() == INTEGER) {
				return GenericRow.of(rnd.nextInt());
			} else {
				return GenericRow.of(GenericRow.of(rnd.nextInt()));
			}
		case ANY:
			return new BinaryGeneric<>(rnd.nextInt(), IntSerializer.INSTANCE);
		default:
			throw new RuntimeException("Not support!");
	}
}
 

private static RowType createJsonRowType(DataType databaseSchema) {
	// Canal JSON contains other information, e.g. "database", "ts"
	// but we don't need them
	return (RowType) DataTypes.ROW(
		DataTypes.FIELD("data", DataTypes.ARRAY(databaseSchema)),
		DataTypes.FIELD("old", DataTypes.ARRAY(databaseSchema)),
		DataTypes.FIELD("type", DataTypes.STRING())).getLogicalType();
}
 

/**
 * RowType of table that excludes partition key fields.
 */
default RowType getFormatRowType() {
	return RowType.of(
		Arrays.stream(getFormatFieldTypes())
			.map(DataType::getLogicalType)
			.toArray(LogicalType[]::new),
		getFormatFieldNames());
}
 

public ArrowPythonScalarFunctionFlatMap(
	Configuration config,
	PythonFunctionInfo[] scalarFunctions,
	RowType inputType,
	RowType outputType,
	int[] udfInputOffsets,
	int[] forwardedFields) {
	super(config, scalarFunctions, inputType, outputType, udfInputOffsets, forwardedFields);
}
 

@Override
public TypeSerializer visit(RowType rowType) {
	final TypeSerializer[] fieldTypeSerializers = rowType.getFields()
		.stream()
		.map(f -> f.getType().accept(this))
		.toArray(TypeSerializer[]::new);
	return new RowDataSerializer(rowType.getChildren().toArray(new LogicalType[0]), fieldTypeSerializers);
}
 

private Function<Object, Object> newStructConverter(FieldsDataType dataType, Schema avroStruct) throws SchemaUtils.IncompatibleSchemaException {
    if (avroStruct.getType() != RECORD ||
            avroStruct.getFields().size() != dataType.getFieldDataTypes().size()) {
        throw new SchemaUtils.IncompatibleSchemaException(
                String.format("Cannot convert Flink type %s to Avro type %s.", dataType.toString(), avroStruct.toString(true)));
    }

    Map<String, DataType> fieldsType = dataType.getFieldDataTypes();
    List<RowType.RowField> fields = ((RowType) dataType.getLogicalType()).getFields();

    List<BiFunction<PositionedGetter, Integer, Object>> fieldConverters = new ArrayList<>();

    for (int i = 0; i < fields.size(); i++) {
        RowType.RowField rf = fields.get(i);
        DataType dt = fieldsType.get(rf.getName());
        Schema.Field at = avroStruct.getFields().get(i);
        fieldConverters.add(newConverter(dt, resolveNullableType(at.schema(), dt.getLogicalType().isNullable())));
    }
    int numFields = fieldsType.size();

    return row -> {
        GenericSchema<GenericRecord> pSchema = SchemaUtils.avroSchema2PulsarSchema(avroStruct);
        GenericRecordBuilder builder = pSchema.newRecordBuilder();
        Row rowX = (Row) row;

        for (int i = 0; i < numFields; i++) {
            if (rowX.getField(i) == null) {
                builder.set(pSchema.getFields().get(i), null);
            } else {
                builder.set(pSchema.getFields().get(i), fieldConverters.get(i).apply(new PositionedGetter(rowX), i));
            }
        }
        return (GenericAvroRecord) builder.build();
    };

}