snuyanzin commented on code in PR #25834:
URL: https://github.com/apache/flink/pull/25834#discussion_r1933526803
##########
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<>();
Review Comment:
```suggestion
Map<String, String> properties = new HashMap<>();
```
nit
--
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]
