@Test
public void testPredictCol() throws Exception {

	TableSchema dataSchema = new TableSchema(
		new String[] {"vec"},
		new TypeInformation<?>[] {Types.STRING}
	);
	Params params = new Params()
		.set(NaiveBayesTextPredictParams.VECTOR_COL, "vec")
		.set(NaiveBayesTextPredictParams.PREDICTION_COL, "pred")
		.set(NaiveBayesTextPredictParams.RESERVED_COLS, new String[] {});
	NaiveBayesTextModelMapper mapper = new NaiveBayesTextModelMapper(modelSchema, dataSchema, params);
	mapper.loadModel(model);
	assertEquals(mapper.map(Row.of("1.0, 1.0, 0.0, 1.0")).getField(0), 1);
	assertEquals(mapper.getOutputSchema(), new TableSchema(new String[] {"pred"},
		new TypeInformation<?>[] {Types.INT}));
}