Skip to content

Commit 2df9bb0

Browse files
timothyw553timothyw553
committed
spark: wire DRC createTable
1 parent c6109f3 commit 2df9bb0

8 files changed

Lines changed: 469 additions & 38 deletions

File tree

‎spark/src/main/scala/org/apache/spark/sql/delta/catalog/AbstractDeltaCatalog.scala‎

Lines changed: 24 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,6 @@ class AbstractDeltaCatalog extends DelegatingCatalogExtension
180180
base
181181
}
182182
}
183-
var locUriOpt = location.map(CatalogUtils.stringToURI)
184183
val existingTableOpt = getExistingTableIfExists(id, Some(ident), operation)
185184
// PROP_IS_MANAGED_LOCATION indicates that the table location is not user-specified but
186185
// system-generated. The table should be created as managed table in this case.
@@ -195,10 +194,23 @@ class AbstractDeltaCatalog extends DelegatingCatalogExtension
195194
} else {
196195
CatalogTableType.EXTERNAL
197196
}
197+
val deltaRestCreate = if (isUnityCatalog && existingTableOpt.isEmpty) {
198+
deltaCatalogClient.prepareCreateTable(
199+
ident,
200+
tableType,
201+
location.map(CatalogUtils.stringToURI))
202+
} else {
203+
None
204+
}
205+
val locUriOpt = deltaRestCreate.map(_.location).orElse(location.map(CatalogUtils.stringToURI))
206+
val tablePropertiesWithDeltaRest =
207+
tableProperties ++ deltaRestCreate.map(_.tableProperties).getOrElse(Map.empty)
208+
val writeOptionsWithDeltaRest =
209+
writeOptions ++ deltaRestCreate.map(_.storageProperties).getOrElse(Map.empty)
198210
val loc = locUriOpt
199211
.orElse(existingTableOpt.flatMap(_.storage.locationUri))
200212
.getOrElse(spark.sessionState.catalog.defaultTablePath(id))
201-
val storage = DataSource.buildStorageFormatFromOptions(writeOptions)
213+
val storage = DataSource.buildStorageFormatFromOptions(writeOptionsWithDeltaRest)
202214
.copy(locationUri = Option(loc))
203215
val commentOpt = Option(allTableProperties.get("comment"))
204216

@@ -211,7 +223,7 @@ class AbstractDeltaCatalog extends DelegatingCatalogExtension
211223
provider = Some(DeltaSourceUtils.ALT_NAME),
212224
partitionColumnNames = newPartitionColumns,
213225
bucketSpec = newBucketSpec,
214-
properties = tableProperties,
226+
properties = tablePropertiesWithDeltaRest,
215227
comment = commentOpt
216228
)
217229

@@ -225,7 +237,7 @@ class AbstractDeltaCatalog extends DelegatingCatalogExtension
225237
val writer = sourceQuery.map { df =>
226238
val catalogTbl = Some(tableDesc)
227239
// For safety, only extract the file system options here, to create deltaLog.
228-
val fileSystemOptions = writeOptions.filter { case (k, _) =>
240+
val fileSystemOptions = writeOptionsWithDeltaRest.filter { case (k, _) =>
229241
DeltaTableUtils.validDeltaTableHadoopPrefixes.exists(k.startsWith)
230242
}
231243
val deltaOptions = new DeltaOptions(
@@ -279,9 +291,14 @@ class AbstractDeltaCatalog extends DelegatingCatalogExtension
279291
// Before this bug is fixed, we should only call the catalog plugin API to create tables
280292
// if UC is enabled to replace `V2SessionCatalog`.
281293
createTableFunc = Option.when(isUnityCatalog) {
282-
v1Table => {
283-
val t = V1Table(v1Table)
284-
super.createTable(ident, t.columns(), t.partitioning, t.properties)
294+
(v1Table, snapshot) => {
295+
deltaRestCreate match {
296+
case Some(_) =>
297+
deltaCatalogClient.createTable(ident, v1Table, snapshot)
298+
case None =>
299+
val t = V1Table(v1Table)
300+
super.createTable(ident, t.columns(), t.partitioning, t.properties)
301+
}
285302
}
286303
}).run(spark)
287304

‎spark/src/main/scala/org/apache/spark/sql/delta/catalog/DeltaCatalogClient.scala‎

Lines changed: 172 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -28,21 +28,43 @@ import io.delta.storage.uc.UCDeltaClient
2828
import io.unitycatalog.client.ApiException
2929
import io.unitycatalog.client.auth.TokenProvider
3030
import io.unitycatalog.client.delta.model.{
31+
CreateTableRequest,
3132
CredentialOperation,
3233
CredentialsResponse,
3334
DataSourceFormat => DeltaDataSourceFormat,
35+
DeltaProtocol => DeltaRestProtocol,
3436
StorageCredential,
37+
StagingTableResponse,
38+
StagingTableResponseRequiredProtocol,
3539
TableType => DeltaTableType
3640
}
3741

3842
import org.apache.spark.sql.SparkSession
3943
import org.apache.spark.sql.catalyst.analysis.NoSuchTableException
4044
import org.apache.spark.sql.catalyst.TableIdentifier
41-
import org.apache.spark.sql.catalyst.catalog.{CatalogStorageFormat, CatalogTable, CatalogTableType, CatalogUtils}
42-
import org.apache.spark.sql.connector.catalog.{CatalogPlugin, Identifier, Table, TableCatalog, V1Table}
45+
import org.apache.spark.sql.catalyst.catalog.{
46+
CatalogStorageFormat,
47+
CatalogTable,
48+
CatalogTableType,
49+
CatalogUtils
50+
}
51+
import org.apache.spark.sql.connector.catalog.{
52+
CatalogPlugin,
53+
Identifier,
54+
Table,
55+
TableCatalog,
56+
V1Table
57+
}
58+
import org.apache.spark.sql.delta.Snapshot
59+
import org.apache.spark.sql.delta.actions.Protocol
4360
import org.apache.spark.sql.delta.coordinatedcommits.UCCommitCoordinatorBuilder
4461
import org.apache.spark.sql.delta.sources.DeltaSourceUtils
4562

63+
private[catalog] case class PreparedDeltaRestCreate(
64+
location: URI,
65+
tableProperties: Map[String, String],
66+
storageProperties: Map[String, String])
67+
4668
private class DeltaCatalogClient private (
4769
private val ucDeltaClient: Option[UCDeltaClient],
4870
delegate: TableCatalog,
@@ -80,6 +102,61 @@ private class DeltaCatalogClient private (
80102
}
81103
}
82104

105+
def prepareCreateTable(
106+
ident: Identifier,
107+
tableType: CatalogTableType,
108+
location: Option[URI]): Option[PreparedDeltaRestCreate] = {
109+
ucDeltaClient match {
110+
case Some(client) if ident.namespace().length == 1 =>
111+
val schemaName = ident.namespace().head
112+
val tableName = ident.name()
113+
(tableType, location) match {
114+
case (CatalogTableType.MANAGED, None) =>
115+
val staging = client.createStagingTable(catalogName, schemaName, tableName)
116+
val stagingLocation = CatalogUtils.stringToURI(staging.getLocation)
117+
Some(PreparedDeltaRestCreate(
118+
location = stagingLocation,
119+
tableProperties = toTableProperties(staging),
120+
storageProperties = toCredentialProperties(
121+
staging.getLocation,
122+
Option(staging.getStorageCredentials).map(_.asScala.toSeq).getOrElse(Nil),
123+
stagingLocation.getScheme)))
124+
case (CatalogTableType.EXTERNAL, Some(locationUri))
125+
if isCloudScheme(locationUri.getScheme) =>
126+
val locationString = locationUri.toString
127+
val credentials = client.getTemporaryPathCredentials(
128+
locationString,
129+
CredentialOperation.READ_WRITE)
130+
Some(PreparedDeltaRestCreate(
131+
location = locationUri,
132+
tableProperties = Map.empty,
133+
storageProperties = toCredentialProperties(
134+
locationString,
135+
getStorageCredentials(credentials),
136+
locationUri.getScheme)))
137+
case _ =>
138+
None
139+
}
140+
case _ =>
141+
None
142+
}
143+
}
144+
145+
def createTable(
146+
ident: Identifier,
147+
table: CatalogTable,
148+
snapshot: Snapshot): Unit = {
149+
ucDeltaClient match {
150+
case Some(client) if ident.namespace().length == 1 =>
151+
client.createTable(
152+
catalogName,
153+
ident.namespace().head,
154+
toCreateTableRequest(ident, table, snapshot))
155+
case _ =>
156+
throw new IllegalStateException(s"Delta REST createTable is not available for $ident.")
157+
}
158+
}
159+
83160
private def translateLoadTableException(ident: Identifier, e: IOException): Throwable = {
84161
e.getCause match {
85162
case api: ApiException if api.getCode == 404 =>
@@ -128,25 +205,94 @@ private class DeltaCatalogClient private (
128205
metadata: io.unitycatalog.client.delta.model.TableMetadata,
129206
credentials: Option[CredentialsResponse],
130207
locationScheme: String): Map[String, String] = {
131-
val storageCredentials = credentials.toSeq.flatMap(getStorageCredentials)
132-
val credentialProperties =
133-
if (!isCloudScheme(locationScheme)) {
134-
Map.empty[String, String]
135-
} else if (storageCredentials.isEmpty) {
136-
throw new IllegalArgumentException(
137-
s"Delta REST returned no storage credentials for cloud location ${metadata.getLocation}.")
138-
} else {
139-
selectStorageCredential(metadata.getLocation, storageCredentials)
140-
.map(storageCredentialToProperties)
141-
.map(withOptionPrefix)
142-
.getOrElse {
143-
throw new IllegalArgumentException(
144-
s"No storage credential matched Delta REST location ${metadata.getLocation}.")
145-
}
146-
}
208+
val credentialProperties = withOptionPrefix(
209+
toCredentialProperties(
210+
metadata.getLocation,
211+
credentials.toSeq.flatMap(getStorageCredentials),
212+
locationScheme))
147213
Map(UC_TABLE_ID_KEY -> metadata.getTableUuid.toString) ++ credentialProperties
148214
}
149215

216+
private def toCredentialProperties(
217+
location: String,
218+
storageCredentials: Seq[StorageCredential],
219+
locationScheme: String): Map[String, String] = {
220+
if (!isCloudScheme(locationScheme)) {
221+
Map.empty[String, String]
222+
} else if (storageCredentials.isEmpty) {
223+
throw new IllegalArgumentException(
224+
s"Delta REST returned no storage credentials for cloud location $location.")
225+
} else {
226+
selectStorageCredential(location, storageCredentials)
227+
.map(storageCredentialToProperties)
228+
.getOrElse {
229+
throw new IllegalArgumentException(
230+
s"No storage credential matched Delta REST location $location.")
231+
}
232+
}
233+
}
234+
235+
private def toTableProperties(staging: StagingTableResponse): Map[String, String] = {
236+
protocolFeatureProperties(staging.getRequiredProtocol) ++
237+
Option(staging.getRequiredProperties)
238+
.map(_.asScala.collect { case (key, value) if value != null => key -> value }.toMap)
239+
.getOrElse(Map.empty) ++
240+
Map(
241+
TableCatalog.PROP_IS_MANAGED_LOCATION -> "true",
242+
UC_TABLE_ID_KEY -> staging.getTableId.toString)
243+
}
244+
245+
private def protocolFeatureProperties(
246+
protocol: StagingTableResponseRequiredProtocol): Map[String, String] = {
247+
Option(protocol).map { p =>
248+
(Option(p.getReaderFeatures).map(_.asScala).getOrElse(Nil) ++
249+
Option(p.getWriterFeatures).map(_.asScala).getOrElse(Nil))
250+
.map(feature => s"delta.feature.$feature" -> "supported")
251+
.toMap
252+
}.getOrElse(Map.empty)
253+
}
254+
255+
private def toCreateTableRequest(
256+
ident: Identifier,
257+
table: CatalogTable,
258+
snapshot: Snapshot): CreateTableRequest = {
259+
new CreateTableRequest()
260+
.name(ident.name())
261+
.location(table.storage.locationUri
262+
.getOrElse {
263+
throw new IllegalArgumentException(
264+
s"Delta REST createTable requires a location for ${ident.toString}.")
265+
}
266+
.toString)
267+
.tableType(toDeltaTableType(table.tableType))
268+
.dataSourceFormat(DeltaDataSourceFormat.DELTA)
269+
.comment(table.comment.orNull)
270+
.columns(DeltaRestSchemaConverter.toDeltaType(snapshot.schema))
271+
.partitionColumns(snapshot.metadata.partitionColumns.asJava)
272+
.protocol(toDeltaProtocol(snapshot.protocol))
273+
.properties(toDeltaCreateTableProperties(snapshot.metadata.configuration).asJava)
274+
}
275+
276+
private def toDeltaTableType(tableType: CatalogTableType): DeltaTableType = tableType match {
277+
case CatalogTableType.MANAGED => DeltaTableType.MANAGED
278+
case CatalogTableType.EXTERNAL => DeltaTableType.EXTERNAL
279+
case other =>
280+
throw new IllegalArgumentException(s"Unsupported Delta REST table type: $other")
281+
}
282+
283+
private def toDeltaProtocol(protocol: Protocol): DeltaRestProtocol = {
284+
new DeltaRestProtocol()
285+
.minReaderVersion(protocol.minReaderVersion)
286+
.minWriterVersion(protocol.minWriterVersion)
287+
.readerFeatures(protocol.readerFeatureNames.toSeq.sorted.asJava)
288+
.writerFeatures(protocol.writerFeatureNames.toSeq.sorted.asJava)
289+
}
290+
291+
private def toDeltaCreateTableProperties(
292+
properties: Map[String, String]): Map[String, String] = {
293+
properties -- DeltaCatalogClient.V2CreateTableProperties
294+
}
295+
150296
private def getStorageCredentials(credentials: CredentialsResponse): Seq[StorageCredential] = {
151297
Option(credentials)
152298
.flatMap(c => Option(c.getStorageCredentials))
@@ -219,6 +365,14 @@ private class DeltaCatalogClient private (
219365

220366
private object DeltaCatalogClient {
221367
private val CloudSchemes = Set("s3", "s3a", "gs", "abfs", "abfss")
368+
private val V2CreateTableProperties = Set(
369+
TableCatalog.PROP_COMMENT,
370+
TableCatalog.PROP_EXTERNAL,
371+
TableCatalog.PROP_IS_MANAGED_LOCATION,
372+
TableCatalog.PROP_LOCATION,
373+
TableCatalog.PROP_OWNER,
374+
TableCatalog.PROP_PROVIDER,
375+
"path")
222376

223377
private def isCloudScheme(scheme: String): Boolean = {
224378
Option(scheme).exists(s => CloudSchemes.contains(s.toLowerCase(Locale.ROOT)))

‎spark/src/main/scala/org/apache/spark/sql/delta/catalog/DeltaRestSchemaConverter.scala‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,12 +16,14 @@
1616

1717
package org.apache.spark.sql.delta.catalog
1818

19-
import java.util.{List => JList, Map => JMap}
19+
import java.util.{Collections, List => JList, Map => JMap}
2020

2121
import scala.collection.JavaConverters._
2222

23+
import com.fasterxml.jackson.core.`type`.TypeReference
2324
import io.unitycatalog.client.delta.model
2425

26+
import org.apache.spark.sql.delta.util.JsonUtils
2527
import org.apache.spark.sql.types.{
2628
ArrayType,
2729
DataType,
@@ -34,11 +36,53 @@ import org.apache.spark.sql.types.{
3436
}
3537

3638
object DeltaRestSchemaConverter {
39+
private val MetadataMapType = new TypeReference[JMap[String, Object]] {}
3740

3841
def toSparkType(schema: model.StructType): StructType = {
3942
StructType(schema.getFields.asScala.map(toSparkField).toSeq)
4043
}
4144

45+
def toDeltaType(schema: StructType): model.StructType = {
46+
new model.StructType()
47+
.fields(schema.fields.map(toDeltaField).toSeq.asJava)
48+
}
49+
50+
private def toDeltaField(field: StructField): model.StructField = {
51+
new model.StructField()
52+
.name(field.name)
53+
.`type`(toDeltaType(field.dataType))
54+
.nullable(field.nullable)
55+
.metadata(toDeltaMetadata(field.metadata))
56+
}
57+
58+
private def toDeltaType(dataType: DataType): model.DeltaType = dataType match {
59+
case struct: StructType =>
60+
toDeltaType(struct)
61+
case ArrayType(elementType, containsNull) =>
62+
new model.ArrayType()
63+
.elementType(toDeltaType(elementType))
64+
.containsNull(containsNull)
65+
case MapType(keyType, valueType, valueContainsNull) =>
66+
new model.MapType()
67+
.keyType(toDeltaType(keyType))
68+
.valueType(toDeltaType(valueType))
69+
.valueContainsNull(valueContainsNull)
70+
case decimal: DecimalType =>
71+
new model.DecimalType()
72+
.precision(decimal.precision)
73+
.scale(decimal.scale)
74+
case primitive =>
75+
new model.PrimitiveType().`type`(primitive.typeName)
76+
}
77+
78+
private def toDeltaMetadata(metadata: Metadata): JMap[String, Object] = {
79+
if (metadata == null || metadata.isEmpty) {
80+
Collections.emptyMap()
81+
} else {
82+
JsonUtils.mapper.readValue(metadata.json, MetadataMapType)
83+
}
84+
}
85+
4286
private def toSparkField(field: model.StructField): StructField = {
4387
StructField(
4488
name = field.getName,

‎spark/src/main/scala/org/apache/spark/sql/delta/commands/CreateDeltaTableCommand.scala‎

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -67,8 +67,9 @@ import org.apache.spark.util.Utils
6767
* @param output SQL output of the command
6868
* @param protocol This is used to create a table with specific protocol version
6969
* @param allowCatalogManaged This is used to create UC managed table with catalogManaged feature
70-
* @param createTableFunc If specified, call this function to create the table, instead of
71-
* Spark `SessionCatalog#createTable` which is backed by Hive Metastore.
70+
* @param createTableFunc If specified, call this function with the post-commit snapshot to create
71+
* the table, instead of Spark `SessionCatalog#createTable` which is backed
72+
* by Hive Metastore.
7273
*/
7374
case class CreateDeltaTableCommand(
7475
override val table: CatalogTable,
@@ -80,7 +81,7 @@ case class CreateDeltaTableCommand(
8081
override val output: Seq[Attribute] = Nil,
8182
protocol: Option[Protocol] = None,
8283
override val allowCatalogManaged: Boolean = false,
83-
createTableFunc: Option[CatalogTable => Unit] = None)
84+
createTableFunc: Option[(CatalogTable, Snapshot) => Unit] = None)
8485
extends LeafRunnableCommand
8586
with DeltaCommand
8687
with DeltaLogging

0 commit comments

Comments
 (0)