Skip to content

[SPARK-37343][SQL] Implement createIndex, IndexExists and dropIndex in JDBC (Postgres dialect) #34673

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
dchvn wants to merge 6 commits into from
Closed

[SPARK-37343][SQL] Implement createIndex, IndexExists and dropIndex in JDBC (Postgres dialect) #34673

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
baae935
Dsv2_index_postgres
dchvn Nov 12, 2021
0e04370
move code to JdbcUtils
dchvn Nov 30, 2021
ac56e68
update for index type postgresql
dchvn Dec 1, 2021
2086915
resolve review
dchvn Dec 2, 2021
67338d4
update code
dchvn Dec 7, 2021
36ee897
resolve review
dchvn Dec 16, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
46 changes: 36 additions & 10 deletions sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/jdbc/JdbcUtils.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ import java.util
import java.util.Locale
import java.util.concurrent.TimeUnit

import scala.collection.JavaConverters._
import scala.util.Try
import scala.util.control.NonFatal

Expand All @@ -38,6 +39,7 @@ import org.apache.spark.sql.catalyst.parser.CatalystSqlParser
import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateTimeUtils, GenericArrayData}
import org.apache.spark.sql.catalyst.util.DateTimeUtils.{instantToMicros, localDateToDays, toJavaDate, toJavaTimestamp}
import org.apache.spark.sql.connector.catalog.TableChange
import org.apache.spark.sql.connector.catalog.index.SupportsIndex
Copy link
Contributor

@huaxingao huaxingao Dec 6, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: combine this with the next line

import org.apache.spark.sql.connector.catalog.index.TableIndex
import org.apache.spark.sql.connector.expressions.NamedReference
import org.apache.spark.sql.errors.{QueryCompilationErrors, QueryExecutionErrors}
Expand Down Expand Up @@ -1025,7 +1027,7 @@ object JdbcUtils extends Logging with SQLConfHelper {
options: JDBCOptions): Unit = {
val dialect = JdbcDialects.get(options.url)
executeStatement(conn, options,
dialect.createIndex(indexName, tableName, columns, columnsProperties, properties))
dialect.createIndex(indexName, tableName, columns, columnsProperties, properties, options))
}

/**
Expand Down Expand Up @@ -1078,26 +1080,50 @@ object JdbcUtils extends Logging with SQLConfHelper {
*/
def checkIfIndexExists(
conn: Connection,
indexName: String,
sql: String,
indexColumnName: String,
options: JDBCOptions): Boolean = {
val statement = conn.createStatement
try {
statement.setQueryTimeout(options.queryTimeout)
val rs = statement.executeQuery(sql)
while (rs.next()) {
val retrievedIndexName = rs.getString(indexColumnName)
if (conf.resolver(retrievedIndexName, indexName)) {
return true
}
}
false
rs.next
} catch {
case _: Exception =>
logWarning("Cannot retrieved index info.")
false
} finally {
statement.close()
}
}

/**
* Process index properties and return tuple of indexType and list of the other index properties.
*/
def processIndexProperties(
properties: util.Map[String, String],
options: JDBCOptions
): (String, Array[String]) = {
val dialect = JdbcDialects.get(options.url)
var indexType = ""
var indexPropertyList: Array[String] = Array.empty
Copy link
Contributor

@cloud-fan cloud-fan Dec 1, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this should be ArrayBuffer


if (!properties.isEmpty) {
properties.asScala.foreach { case (k, v) =>
if (k.equals(SupportsIndex.PROP_TYPE)) {
if (v.equalsIgnoreCase("BTREE") || v.equalsIgnoreCase("HASH")) {
indexType = s"USING $v"
} else {
throw new UnsupportedOperationException(s"Index Type $v is not supported." +
" The supported Index Types are: BTREE and HASH")
Copy link
Contributor

@huaxingao huaxingao Nov 30, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Seems PostgreSQL can support more index types (https://www.postgresql.org/docs/9.5/sql-createindex.html).

Copy link
Contributor Author

@dchvn dchvn Dec 1, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@huaxingao, GIST and GIN index types in PostgreSQL are used for specific types: tsvector and tsquery, which are not supported types in Spark. Also for SP-GIST are used for the point type. Should we ignore these index types or something else? GIST and GIN index type for details.

}
} else {
indexPropertyList = indexPropertyList :+ dialect.convertPropertyPairToString(k, v)
}
}
}
(indexType, indexPropertyList)
}
Copy link
Contributor

@huaxingao huaxingao Dec 6, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think we don't need this method for now. Seems to me that databases always allow this = in between of key and value. I didn't verify this syntax in all the databases, but I checked MySQL, it takes = as optional.


def executeQuery(conn: Connection, options: JDBCOptions, sql: String): ResultSet = {
val statement = conn.createStatement
try {
Expand Down
10 changes: 9 additions & 1 deletion sql/core/src/main/scala/org/apache/spark/sql/jdbc/JdbcDialects.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,8 @@ abstract class JdbcDialect extends Serializable with Logging{
tableName: String,
columns: Array[NamedReference],
columnsProperties: util.Map[NamedReference, util.Map[String, String]],
properties: util.Map[String, String]): String = {
properties: util.Map[String, String],
options: JDBCOptions): String = {
throw new UnsupportedOperationException("createIndex is not supported")
}

Expand Down Expand Up @@ -358,6 +359,13 @@ abstract class JdbcDialect extends Serializable with Logging{
new AnalysisException(message, cause = Some(e))
}

/**
* Convert key-value property pair to string
*/
def convertPropertyPairToString(key: String, value: String): String = {
Copy link
Contributor

@cloud-fan cloud-fan Dec 1, 2021

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note that, every method in this class is a public API, and we should be very careful when adding new public APIs.

I don't think the two newly added APIs are necessary. They are just used to share the code between different dialects. We should use internal util functions to share code.

s"$key $value"
}

/**
* returns the LIMIT clause for the SELECT statement
*/
Expand Down
38 changes: 9 additions & 29 deletions sql/core/src/main/scala/org/apache/spark/sql/jdbc/MySQLDialect.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,10 @@ import java.sql.{Connection, SQLException, Types}
import java.util
import java.util.Locale

import scala.collection.JavaConverters._

import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.SQLConfHelper
import org.apache.spark.sql.catalyst.analysis.{IndexAlreadyExistsException, NoSuchIndexException}
import org.apache.spark.sql.connector.catalog.index.{SupportsIndex, TableIndex}
import org.apache.spark.sql.connector.catalog.index.{TableIndex}
import org.apache.spark.sql.connector.expressions.{FieldReference, NamedReference}
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.datasources.jdbc.{JDBCOptions, JdbcUtils}
Expand Down Expand Up @@ -118,27 +116,15 @@ private case object MySQLDialect extends JdbcDialect with SQLConfHelper {
tableName: String,
columns: Array[NamedReference],
columnsProperties: util.Map[NamedReference, util.Map[String, String]],
properties: util.Map[String, String]): String = {
properties: util.Map[String, String],
options: JDBCOptions): String = {
val columnList = columns.map(col => quoteIdentifier(col.fieldNames.head))
var indexProperties: String = ""
var indexType = ""
if (!properties.isEmpty) {
properties.asScala.foreach { case (k, v) =>
if (k.equals(SupportsIndex.PROP_TYPE)) {
if (v.equalsIgnoreCase("BTREE") || v.equalsIgnoreCase("HASH")) {
indexType = s"USING $v"
} else {
throw new UnsupportedOperationException(s"Index Type $v is not supported." +
" The supported Index Types are: BTREE and HASH")
}
} else {
indexProperties = indexProperties + " " + s"$k $v"
}
}
}
val (indexType, indexPropertyList) = JdbcUtils.processIndexProperties(properties, options)

// columnsProperties doesn't apply to MySQL so it is ignored
s"CREATE INDEX ${quoteIdentifier(indexName)} $indexType ON" +
s" ${quoteIdentifier(tableName)} (${columnList.mkString(", ")}) $indexProperties"
s" ${quoteIdentifier(tableName)} (${columnList.mkString(", ")})" +
s" ${indexPropertyList.mkString(" ")}"
}

// SHOW INDEX syntax
Expand All @@ -148,14 +134,8 @@ private case object MySQLDialect extends JdbcDialect with SQLConfHelper {
indexName: String,
tableName: String,
options: JDBCOptions): Boolean = {
val sql = s"SHOW INDEXES FROM ${quoteIdentifier(tableName)}"
try {
JdbcUtils.checkIfIndexExists(conn, indexName, sql, "key_name", options)
} catch {
case _: Exception =>
logWarning("Cannot retrieved index info.")
false
}
val sql = s"SHOW INDEXES FROM ${quoteIdentifier(tableName)} WHERE key_name = '$indexName'"
JdbcUtils.checkIfIndexExists(conn, sql, options)
}

override def dropIndex(indexName: String, tableName: String): String = {
Expand Down
50 changes: 15 additions & 35 deletions sql/core/src/main/scala/org/apache/spark/sql/jdbc/PostgresDialect.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,12 +21,9 @@ import java.sql.{Connection, SQLException, Types}
import java.util
import java.util.Locale

import scala.collection.JavaConverters._

import org.apache.spark.sql.AnalysisException
import org.apache.spark.sql.catalyst.SQLConfHelper
import org.apache.spark.sql.catalyst.analysis.{IndexAlreadyExistsException, NoSuchIndexException}
import org.apache.spark.sql.connector.catalog.index.SupportsIndex
import org.apache.spark.sql.connector.expressions.NamedReference
import org.apache.spark.sql.execution.datasources.jdbc.{JDBCOptions, JdbcUtils}
import org.apache.spark.sql.execution.datasources.v2.TableSampleInfo
Expand Down Expand Up @@ -180,34 +177,18 @@ private object PostgresDialect extends JdbcDialect with SQLConfHelper {
tableName: String,
columns: Array[NamedReference],
columnsProperties: util.Map[NamedReference, util.Map[String, String]],
properties: util.Map[String, String]): String = {
properties: util.Map[String, String],
options: JDBCOptions): String = {
val columnList = columns.map(col => quoteIdentifier(col.fieldNames.head))
var indexPropertiesStr: String = ""
var hasIndexProperties: Boolean = false
var indexType = ""

if (!properties.isEmpty) {
var indexPropertyList: Array[String] = Array.empty
properties.asScala.foreach { case (k, v) =>
if (k.equals(SupportsIndex.PROP_TYPE)) {
if (v.equalsIgnoreCase("BTREE") || v.equalsIgnoreCase("HASH")) {
indexType = s"USING $v"
} else {
throw new UnsupportedOperationException(s"Index Type $v is not supported." +
" The supported Index Types are: BTREE and HASH")
}
} else {
hasIndexProperties = true
indexPropertyList = indexPropertyList :+ s"$k = $v"
}
}
if (hasIndexProperties) {
indexPropertiesStr += "WITH (" + indexPropertyList.mkString(", ") + ")"
}
var indexProperties = ""
val (indexType, indexPropertyList) = JdbcUtils.processIndexProperties(properties, options)

if (indexPropertyList.nonEmpty) {
indexProperties = "WITH (" + indexPropertyList.mkString(", ") + ")"
}

s"CREATE INDEX ${quoteIdentifier(indexName)} ON ${quoteIdentifier(tableName)}" +
s" $indexType (${columnList.mkString(", ")}) $indexPropertiesStr"
s" $indexType (${columnList.mkString(", ")}) $indexProperties"
}

// SHOW INDEX syntax
Expand All @@ -217,14 +198,9 @@ private object PostgresDialect extends JdbcDialect with SQLConfHelper {
indexName: String,
tableName: String,
options: JDBCOptions): Boolean = {
val sql = s"SELECT * FROM pg_indexes WHERE tablename = '$tableName'"
try {
JdbcUtils.checkIfIndexExists(conn, indexName, sql, "indexname", options)
} catch {
case _: Exception =>
logWarning("Cannot retrieved index info.")
false
}
val sql = s"SELECT * FROM pg_indexes WHERE tablename = '$tableName' AND" +
s" indexname = '$indexName'"
JdbcUtils.checkIfIndexExists(conn, sql, options)
}

// DROP INDEX syntax
Expand All @@ -246,4 +222,8 @@ private object PostgresDialect extends JdbcDialect with SQLConfHelper {
case _ => super.classifyException(message, e)
}
}

override def convertPropertyPairToString(key: String, value: String): String = {
s"$key = $value"
}
}