chl-wxp commented on code in PR #10988:
URL: https://github.com/apache/seatunnel/pull/10988#discussion_r3353023587
##########
seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/dialect/psql/PostgresTypeMapper.java:
##########
@@ -51,4 +59,107 @@ public Column mappingColumn(ResultSetMetaData metadata, int
colIndex) throws SQL
.build();
return mappingColumn(typeDefine);
}
+
+ /**
+ * Overrides the default mapping to resolve vector dimension from
pg_attribute when the column
+ * type is {@code vector}. ResultSetMetaData does not carry the dimension
for pgvector columns,
+ * so we query {@code format_type(atttypid, atttypmod)} to obtain the full
type string (e.g.
+ * "vector(3)").
+ */
+ @Override
+ public Column mappingColumn(
+ ResultSetMetaData metadata, int colIndex, Connection connection,
String sqlQuery)
+ throws SQLException {
+ String nativeType = metadata.getColumnTypeName(colIndex);
+
+ if (PostgresTypeConverter.PG_VECTOR.equalsIgnoreCase(nativeType) &&
connection != null) {
+ String tableName = resolveTableName(metadata, colIndex, sqlQuery);
+ if (tableName != null && !tableName.isEmpty()) {
+ String columnName = metadata.getColumnLabel(colIndex);
Review Comment:
For vector dimension lookup, should we use ResultSetMetaData#getColumnName
instead of getColumnLabel?
The output Column name can still use getColumnLabel, but
pg_attribute.attname should be matched against the physical source column name.
Otherwise a query like `SELECT embedding AS embedding_vector FROM public.items`
may fail to recover vector(N), because pg_attribute contains `embedding`, not
the alias `embedding_vector`.
##########
seatunnel-e2e/seatunnel-connector-v2-e2e/connector-jdbc-e2e/connector-jdbc-e2e-part-3/src/test/java/org/apache/seatunnel/connectors/seatunnel/jdbc/JdbcPostgresPgVectorIT.java:
##########
@@ -0,0 +1,264 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.seatunnel.connectors.seatunnel.jdbc;
+
+import org.apache.seatunnel.e2e.common.TestResource;
+import org.apache.seatunnel.e2e.common.TestSuiteBase;
+import org.apache.seatunnel.e2e.common.container.ContainerExtendedFactory;
+import org.apache.seatunnel.e2e.common.container.EngineType;
+import org.apache.seatunnel.e2e.common.container.TestContainer;
+import org.apache.seatunnel.e2e.common.junit.DisabledOnContainer;
+import org.apache.seatunnel.e2e.common.junit.TestContainerExtension;
+import org.apache.seatunnel.e2e.common.util.JdbcUtil;
+
+import org.junit.jupiter.api.AfterAll;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeAll;
+import org.junit.jupiter.api.TestTemplate;
+import org.testcontainers.containers.Container;
+import org.testcontainers.containers.PostgreSQLContainer;
+import org.testcontainers.containers.output.Slf4jLogConsumer;
+import org.testcontainers.lifecycle.Startables;
+import org.testcontainers.utility.DockerImageName;
+import org.testcontainers.utility.DockerLoggerFactory;
+
+import lombok.extern.slf4j.Slf4j;
+
+import java.io.IOException;
+import java.sql.Connection;
+import java.sql.DriverManager;
+import java.sql.SQLException;
+import java.sql.Statement;
+import java.util.List;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Stream;
+
+import static org.awaitility.Awaitility.given;
+
+@Slf4j
+@DisabledOnContainer(
+ value = {},
+ type = {EngineType.SPARK, EngineType.FLINK},
+ disabledReason = "Currently SPARK and FLINK not support FLOAT_VECTOR
type")
+public class JdbcPostgresPgVectorIT extends TestSuiteBase implements
TestResource {
+
+ private static final String PGVECTOR_IMAGE = "pgvector/pgvector:pg16";
+ private static final String PG_DRIVER_JAR =
+
"https://repo1.maven.org/maven2/org/postgresql/postgresql/42.3.3/postgresql-42.3.3.jar";;
+ private static final String PGVECTOR_SOURCE_DDL =
+ "CREATE EXTENSION IF NOT EXISTS vector;\n"
+ + "CREATE TABLE IF NOT EXISTS pgvector_e2e_source_table
(\n"
+ + " id SERIAL PRIMARY KEY,\n"
+ + " name VARCHAR(255),\n"
+ + " embedding vector(3)\n"
+ + ")";
+ private static final String PGVECTOR_SINK_DDL =
+ "CREATE EXTENSION IF NOT EXISTS vector;\n"
+ + "CREATE TABLE IF NOT EXISTS pgvector_e2e_sink_table (\n"
+ + " id SERIAL PRIMARY KEY,\n"
+ + " name VARCHAR(255),\n"
+ + " embedding vector(3)\n"
+ + ")";
+ private static final String PGVECTOR_AUTO_CREATE_SOURCE_DDL =
+ "CREATE EXTENSION IF NOT EXISTS vector;\n"
+ + "CREATE TABLE IF NOT EXISTS
pgvector_e2e_auto_create_source_table (\n"
+ + " id SERIAL PRIMARY KEY,\n"
+ + " name VARCHAR(255),\n"
+ + " embedding vector(3)\n"
+ + ")";
+
+ private PostgreSQLContainer<?> PGVECTOR_CONTAINER;
+
+ @TestContainerExtension
+ private final ContainerExtendedFactory extendedFactory =
+ container -> {
+ Container.ExecResult extraCommands =
+ container.execInContainer(
+ "bash",
+ "-c",
+ "mkdir -p /tmp/seatunnel/plugins/Jdbc/lib &&
cd /tmp/seatunnel/plugins/Jdbc/lib && curl -O "
+ + PG_DRIVER_JAR);
+ Assertions.assertEquals(0, extraCommands.getExitCode());
+ };
+
+ @BeforeAll
+ @Override
+ public void startUp() throws Exception {
+ PGVECTOR_CONTAINER =
+ new PostgreSQLContainer<>(
+ DockerImageName.parse(PGVECTOR_IMAGE)
+ .asCompatibleSubstituteFor("postgres"))
+ .withNetwork(TestSuiteBase.NETWORK)
+ .withNetworkAliases("pgvector")
+ .withLogConsumer(
+ new Slf4jLogConsumer(
+
DockerLoggerFactory.getLogger(PGVECTOR_IMAGE)));
+ Startables.deepStart(Stream.of(PGVECTOR_CONTAINER)).join();
+ log.info("PgVector container started");
+ Class.forName(PGVECTOR_CONTAINER.getDriverClassName());
+ given().ignoreExceptions()
+ .await()
+ .atLeast(100, TimeUnit.MILLISECONDS)
+ .pollInterval(500, TimeUnit.MILLISECONDS)
+ .atMost(2, TimeUnit.MINUTES)
+ .untilAsserted(this::initializeJdbcTable);
+ log.info("pgvector data initialization succeeded");
+ }
+
+ @TestTemplate
+ public void testPgVectorSourceAndSink(TestContainer container)
+ throws IOException, InterruptedException {
+ Container.ExecResult execResult =
+
container.executeJob("/jdbc_postgres_pgvector_source_and_sink.conf");
+ Assertions.assertEquals(0, execResult.getExitCode(), "pgvector job run
failed");
+
+ List<List<Object>> src =
+ querySql(
+ "select id, name, embedding::text from
pgvector_e2e_source_table order by id");
+ List<List<Object>> dst =
+ querySql(
+ "select id, name, embedding::text from
pgvector_e2e_sink_table order by id");
+ Assertions.assertFalse(src.isEmpty());
+ Assertions.assertEquals(src.size(), dst.size());
+ for (int i = 0; i < src.size(); i++) {
+ Assertions.assertEquals(src.get(i).get(0), dst.get(i).get(0));
+ Assertions.assertEquals(src.get(i).get(1), dst.get(i).get(1));
+ assertVectorEquals(src.get(i).get(2).toString(),
dst.get(i).get(2).toString());
+ }
+ log.info("pgvector e2e test completed");
+ }
+
+ /**
+ * Tests the auto-create sink path: sink table is NOT pre-created, so
SeaTunnel must infer the
+ * schema from the source query (including vector dimension) and generate
the sink DDL
+ * automatically. This validates that vector(N) dimensions survive the
full query-source ->
+ * auto-create-sink pipeline.
+ */
+ @TestTemplate
+ public void testPgVectorAutoCreateSink(TestContainer container)
+ throws IOException, InterruptedException {
+ Container.ExecResult execResult =
+
container.executeJob("/jdbc_postgres_pgvector_auto_create_sink.conf");
+ Assertions.assertEquals(0, execResult.getExitCode(), "pgvector
auto-create job run failed");
+
+ // Verify data was written correctly
+ List<List<Object>> src =
+ querySql(
+ "select id, name, embedding::text from
pgvector_e2e_auto_create_source_table order by id");
+ List<List<Object>> dst =
+ querySql(
+ "select id, name, embedding::text from
pgvector_e2e_auto_create_sink_table order by id");
+ Assertions.assertFalse(src.isEmpty());
+ Assertions.assertEquals(src.size(), dst.size());
+ for (int i = 0; i < src.size(); i++) {
+ Assertions.assertEquals(src.get(i).get(0), dst.get(i).get(0));
+ Assertions.assertEquals(src.get(i).get(1), dst.get(i).get(1));
+ assertVectorEquals(src.get(i).get(2).toString(),
dst.get(i).get(2).toString());
+ }
+
+ // Verify the sink table was created with the correct vector dimension
+ List<List<Object>> columns =
+ querySql(
+ "SELECT column_name, data_type FROM
information_schema.columns "
+ + "WHERE table_name =
'pgvector_e2e_auto_create_sink_table' AND column_name = 'embedding'");
+ Assertions.assertFalse(
+ columns.isEmpty(), "embedding column should exist in
auto-created sink table");
+ // Note: information_schema.data_type may show "USER-DEFINED" for
vector,
+ // so we also check via pg_attribute to confirm the dimension
+ List<List<Object>> vectorInfo =
+ querySql(
+ "SELECT format_type(atttypid, atttypmod) FROM
pg_attribute "
+ + "WHERE attrelid =
'pgvector_e2e_auto_create_sink_table'::regclass AND attname = 'embedding'");
+ Assertions.assertFalse(
+ vectorInfo.isEmpty(), "embedding column should exist in
pg_attribute");
+ String fullType = vectorInfo.get(0).get(0).toString();
+ Assertions.assertTrue(
+ fullType.startsWith("vector(") && fullType.endsWith(")"),
+ "Expected vector(N) type but got: " + fullType);
Review Comment:
check vector(N)
##########
seatunnel-connectors-v2/connector-jdbc/src/main/java/org/apache/seatunnel/connectors/seatunnel/jdbc/internal/dialect/psql/PostgresTypeMapper.java:
##########
@@ -51,4 +59,107 @@ public Column mappingColumn(ResultSetMetaData metadata, int
colIndex) throws SQL
.build();
return mappingColumn(typeDefine);
}
+
+ /**
+ * Overrides the default mapping to resolve vector dimension from
pg_attribute when the column
+ * type is {@code vector}. ResultSetMetaData does not carry the dimension
for pgvector columns,
+ * so we query {@code format_type(atttypid, atttypmod)} to obtain the full
type string (e.g.
+ * "vector(3)").
+ */
+ @Override
+ public Column mappingColumn(
+ ResultSetMetaData metadata, int colIndex, Connection connection,
String sqlQuery)
+ throws SQLException {
+ String nativeType = metadata.getColumnTypeName(colIndex);
+
+ if (PostgresTypeConverter.PG_VECTOR.equalsIgnoreCase(nativeType) &&
connection != null) {
+ String tableName = resolveTableName(metadata, colIndex, sqlQuery);
+ if (tableName != null && !tableName.isEmpty()) {
+ String columnName = metadata.getColumnLabel(colIndex);
+ String vectorType =
+ queryVectorTypeFromPgAttribute(connection, tableName,
columnName);
+ if (vectorType != null) {
+ BasicTypeDefine typeDefine =
+ BasicTypeDefine.builder()
+ .name(columnName)
+ .columnType(vectorType)
+ .dataType(nativeType)
+ .sqlType(metadata.getColumnType(colIndex))
+ .nullable(
+ metadata.isNullable(colIndex)
+ ==
ResultSetMetaData.columnNullable)
+ .length((long)
metadata.getPrecision(colIndex))
+ .precision((long)
metadata.getPrecision(colIndex))
+ .scale(metadata.getScale(colIndex))
+ .build();
+ return mappingColumn(typeDefine);
+ }
+ }
+ }
+
+ return mappingColumn(metadata, colIndex);
+ }
+
+ /**
+ * Resolves the table name by first checking ResultSetMetaData, then
falling back to parsing the
+ * SQL query.
+ */
+ private String resolveTableName(ResultSetMetaData metadata, int colIndex,
String sqlQuery) {
+ // Try ResultSetMetaData first — the PostgreSQL JDBC driver often
populates this
+ try {
+ String tableName = metadata.getTableName(colIndex);
+ if (tableName != null && !tableName.isEmpty()) {
+ return tableName;
+ }
+ } catch (SQLException ignored) {
+ // driver may not support it
+ }
+
+ // Fall back to parsing FROM clause from the query
+ if (sqlQuery != null) {
+ Matcher matcher = FROM_TABLE_PATTERN.matcher(sqlQuery);
+ if (matcher.find()) {
+ String fullName = matcher.group(1);
+ // Strip any quoting and extract the leaf table name
+ fullName = fullName.replace("\"", "").replace("`", "");
+ int lastDot = fullName.lastIndexOf('.');
+ return lastDot >= 0 ? fullName.substring(lastDot + 1) :
fullName;
+ }
+ }
+ return null;
+ }
+
+ /**
+ * Queries pg_attribute to get the full vector type string (e.g.
"vector(3)") for a specific
+ * column.
+ */
+ private String queryVectorTypeFromPgAttribute(
+ Connection connection, String tableName, String columnName) {
+ String sql =
+ "SELECT format_type(a.atttypid, a.atttypmod) "
+ + "FROM pg_attribute a "
+ + "JOIN pg_class c ON a.attrelid = c.oid "
+ + "WHERE c.relname = ? AND a.attname = ? AND a.attnum
> 0";
Review Comment:
Does this query not use schema as a parameter? Can uniqueness be ensured?
--
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]
