snuyanzin commented on code in PR #25834:
URL: https://github.com/apache/flink/pull/25834#discussion_r1933523153
##########
flink-table/flink-table-planner/src/test/java/org/apache/flink/table/planner/operations/SqlDdlToOperationConverterTest.java:
##########
@@ -244,6 +262,233 @@ public void testAlterDatabase() throws Exception {
.isEqualTo(properties);
}
+ @Test
+ public void testCreateModel() {
+ final String sql =
+ "CREATE MODEL model1 \n"
+ + "INPUT(a bigint comment 'column a', b varchar, c
int, d varchar)\n"
+ + "OUTPUT(e bigint, f int)\n"
+ + " with (\n"
+ + " 'task' = 'clustering',\n"
+ + " 'provider' = 'openai',\n"
+ + " 'openai.endpoint' = 'someendpoint'\n"
+ + ")\n";
+ FlinkPlannerImpl planner = getPlannerBySqlDialect(SqlDialect.DEFAULT);
+ final CalciteParser parser = getParserBySqlDialect(SqlDialect.DEFAULT);
+ Operation operation = parse(sql, planner, parser);
+ assertThat(operation).isInstanceOf(CreateModelOperation.class);
+ CreateModelOperation op = (CreateModelOperation) operation;
+ CatalogModel catalogModel = op.getCatalogModel();
+ assertThat(catalogModel.getOptions())
+ .isEqualTo(
+ ImmutableMap.of(
+ "PROVIDER",
+ "openai",
+ "OPENAI.ENDPOINT",
+ "someendpoint",
+ "TASK",
+ "clustering"));
+ Schema inputSchema = catalogModel.getInputSchema();
+ assertNotNull(inputSchema);
+ assertThat(
+ inputSchema.getColumns().stream()
+ .map(Schema.UnresolvedColumn::getName)
+ .collect(Collectors.toList()))
+ .isEqualTo(Arrays.asList("a", "b", "c", "d"));
+ Schema outputSchema = catalogModel.getOutputSchema();
+ assertNotNull(outputSchema);
+ assertThat(
+ outputSchema.getColumns().stream()
+ .map(Schema.UnresolvedColumn::getName)
+ .collect(Collectors.toList()))
+ .isEqualTo(Arrays.asList("e", "f"));
+ }
+
+ @Test
+ public void testDropModel() throws Exception {
+ Catalog catalog = new GenericInMemoryCatalog("default", "default");
+ if (!catalogManager.getCatalog("cat1").isPresent()) {
+ catalogManager.registerCatalog("cat1", catalog);
+ }
+ catalogManager.createDatabase(
+ "cat1", "db1", new CatalogDatabaseImpl(new HashMap<>(), null),
true);
+
+ Schema inputSchema =
+ Schema.newBuilder()
+ .column("a", DataTypes.INT())
+ .column("b", DataTypes.STRING())
+ .build();
+ Schema outputSchema = Schema.newBuilder().column("label",
DataTypes.STRING()).build();
+ HashMap<String, String> properties = new HashMap<>();
+ properties.put("K1", "v1");
+ CatalogModel catalogModel = CatalogModel.of(inputSchema, outputSchema,
properties, null);
+ catalogManager.setCurrentCatalog("cat1");
+ catalogManager.setCurrentDatabase("db1");
+ ObjectIdentifier modelIdentifier = ObjectIdentifier.of("cat1", "db1",
"m1");
+ catalogManager.createModel(catalogModel, modelIdentifier, true);
Review Comment:
this code snippet seems to be same as for `testAlterModel`
how about extract it to a separate method something like `prepareModel` ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
