snuyanzin commented on code in PR #25834:
URL: https://github.com/apache/flink/pull/25834#discussion_r1933527308
##########
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);
+
+ String[] dropModelSqls =
+ new String[] {
+ "DROP MODEL cat1.db1.m1",
+ "DROP MODEL db1.m1",
+ "DROP MODEL m1",
+ "DROP MODEL IF EXISTS m2"
+ };
+ ObjectIdentifier m1Id = ObjectIdentifier.of("cat1", "db1", "m1");
+ ObjectIdentifier m2Id = ObjectIdentifier.of("cat1", "db1", "m2");
+ ObjectIdentifier[] expectedIdentifiers =
+ new ObjectIdentifier[] {m1Id, m1Id, m1Id, m2Id, m2Id};
+ boolean[] expectedIfExists = new boolean[] {false, false, false, true};
+
+ for (int i = 0; i < dropModelSqls.length; i++) {
+ String sql = dropModelSqls[i];
+ Operation operation = parse(sql);
+ assertThat(operation).isInstanceOf(DropModelOperation.class);
+ final DropModelOperation dropModelOperation = (DropModelOperation)
operation;
+
assertThat(dropModelOperation.getModelIdentifier()).isEqualTo(expectedIdentifiers[i]);
+
assertThat(dropModelOperation.isIfExists()).isEqualTo(expectedIfExists[i]);
+ }
+ }
+
+ @Test
+ public void testDescribeModel() {
+ Operation operation = parse("DESCRIBE MODEL m1");
+ assertThat(operation).isInstanceOf(DescribeModelOperation.class);
+ DescribeModelOperation describeModelOperation =
(DescribeModelOperation) operation;
+ assertThat(describeModelOperation.getSqlIdentifier())
+ .isEqualTo(ObjectIdentifier.of("builtin", "default", "m1"));
+ assertThat(describeModelOperation.isExtended()).isFalse();
+
+ operation = parse("DESCRIBE MODEL EXTENDED m1");
+ assertThat(operation).isInstanceOf(DescribeModelOperation.class);
+ describeModelOperation = (DescribeModelOperation) operation;
+ assertThat(describeModelOperation.isExtended()).isTrue();
+ }
+
+ @ParameterizedTest(name = "{index}: {0}")
+ @MethodSource("inputForShowModelsTest")
+ void testShowModels(String sql, ShowModelsOperation expected, String
expectedSummary) {
+ Operation operation = parse(sql);
+
assertThat(operation).isInstanceOf(ShowModelsOperation.class).isEqualTo(expected);
+ assertThat(operation.asSummaryString()).isEqualTo(expectedSummary);
+ }
+
+ private static Stream<Arguments> inputForShowModelsTest() {
+ return Stream.of(
+ Arguments.of(
+ "SHOW MODELS FROM db1",
+ new ShowModelsOperation("builtin", "db1", "FROM",
null),
+ "SHOW MODELS FROM builtin.db1"),
+ Arguments.of(
+ "SHOW MODELS",
+ new ShowModelsOperation("builtin", "db1", null, null),
+ "SHOW MODELS"),
+ Arguments.of(
+ "SHOW MODELS FROM `builtin`.`db1` LIKE '%abc%'",
+ new ShowModelsOperation(
+ "builtin",
+ "db1",
+ "FROM",
+ ShowLikeOperator.of(LikeType.LIKE, "%abc%")),
+ "SHOW MODELS FROM builtin.db1 LIKE '%abc%'"),
+ Arguments.of(
+ "SHOW MODELS FROM `builtin`.`db1` NOT LIKE '%abc%'",
+ new ShowModelsOperation(
+ "builtin",
+ "db1",
+ "FROM",
+ ShowLikeOperator.of(LikeType.NOT_LIKE,
"%abc%")),
+ "SHOW MODELS FROM builtin.db1 NOT LIKE '%abc%'"));
+ }
+
+ @Test
+ public void testAlterModel() 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]
