diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index def2816607..d818a8e071 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -315,6 +315,7 @@ jobs: org.apache.comet.CometIcebergEncryptionSuite org.apache.comet.CometIcebergRewriteActionSuite org.apache.comet.CometIcebergWriteActionSuite + org.apache.comet.CometIcebergWriteDetectionSuite org.apache.comet.iceberg.IcebergReflectionSuite org.apache.comet.csv.CometCsvNativeReadSuite org.apache.comet.CometFuzzTestSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index bac1dcd15b..71fcf334f6 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -131,6 +131,7 @@ jobs: org.apache.comet.CometIcebergEncryptionSuite org.apache.comet.CometIcebergRewriteActionSuite org.apache.comet.CometIcebergWriteActionSuite + org.apache.comet.CometIcebergWriteDetectionSuite org.apache.comet.iceberg.IcebergReflectionSuite org.apache.comet.csv.CometCsvNativeReadSuite org.apache.comet.CometFuzzTestSuite diff --git a/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md index 8ef6b4cc16..39a50cdd5c 100644 --- a/docs/source/user-guide/latest/iceberg-writes.md +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -41,7 +41,8 @@ writes into two operators: Data files are still written by iceberg-java; only the plan shape changes. The split makes the write's input visible to AQE and to Comet's columnar rules, and it is the groundwork for a planned follow-up in which Comet writes the data files natively via -[iceberg-rust](https://github.com/apache/iceberg-rust). +[iceberg-rust](https://github.com/apache/iceberg-rust), tracked in +[#5308](https://github.com/apache/datafusion-comet/issues/5308). ## Configuration @@ -57,6 +58,9 @@ spark.sql.catalog..warehouse=... # Split-operator plan (experimental, off by default) spark.comet.write.iceberg.splitOperator.enabled=true + +# Native-write eligibility detection (experimental, off by default; requires the split plan) +spark.comet.write.iceberg.nativeAcceleration.enabled=true ``` ## Supported operations @@ -95,3 +99,102 @@ The rewrite is skipped — and the write runs through Spark's stock combined ope In every fallback case the write is planned as if Comet were absent; there is no correctness trade-off, only no plan change. + +## Native Parquet write eligibility + +A planned follow-up ([#5308](https://github.com/apache/datafusion-comet/issues/5308)) replaces +the `IcebergWrite` operator's per-task Parquet write with +[iceberg-rust](https://github.com/apache/iceberg-rust). The native writer must produce the same +outcome as iceberg-java — the same Parquet features, statistics, and manifest metadata — so a +write is only eligible when every table property it depends on is one the native path reproduces +exactly. `spark.comet.write.iceberg.nativeAcceleration.enabled` enables this eligibility check; +with the current release the native writer itself is not yet wired in, so every write still runs +through iceberg-java and the check's outcome is reported as a fall-back reason in Comet's +extended EXPLAIN output. + +**Most Iceberg write settings are not supported.** Detection is an allowlist: a write is +eligible only when its entire effective configuration matches the table below, and anything +else — any other write-affecting property, any key added by a future Iceberg version, any +value outside the supported set, any reflection failure while inspecting the write — falls +back to iceberg-java with a reason reported in extended EXPLAIN. Checks run on the effective +configuration: table properties overlaid with `SparkWrite.writeProperties`, which is where +iceberg-java resolves per-write options and `spark.sql.iceberg.*` session overrides. + +A write is eligible only when ALL of the following hold: + +| Setting | Supported values | +| ------------------------------------------------------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------- | +| resolved write format (`write-format` option overlaid on `write.format.default`) | `parquet` | +| `format-version` | `1` or `2` | +| `write.parquet.compression-codec` / `compression-level` / `row-group-size-bytes` / `page-size-bytes` / `page-row-limit` / `dict-size-bytes` | any value (translated to the native writer) | +| `write.parquet.row-group-check-min-record-count` | unset or `100` (the default) | +| `write.parquet.row-group-check-max-record-count` | unset or `10000` (the default) | +| `write.parquet.page-version` | unset or `v1` | +| `write.parquet.shred-variants` | unset or `false` (Spark 4.x / Iceberg 1.11 resolve this into every parquet write) | +| `write.parquet.variant-inference-buffer-size` | any value (only meaningful when shredding, which is gated) | +| `write.parquet.bloom-filter-enabled.column.` | unset or `false` | +| `write.metadata.metrics.default` | unset, `truncate(N)`, or `full` | +| `write.metadata.metrics.column.` | unset, `truncate(N)`, or `full` | +| `write.spark.fanout.enabled` | any value (the native writer implements both clustered and fanout modes) | +| `write.target-file-size-bytes` | any value (file rolling cadence differs; see accepted divergences) | +| data location URI scheme | `file`, `memory`, `s3`, `s3a`, `gs`, `oss` | +| partition spec | any (but see partition paths under accepted divergences) | + +Within the namespaces that shape data-file bytes — `write.parquet.*` and `parquet.*` — +everything not listed above must be absent: unvetted `write.parquet.*` keys (e.g. +`bloom-filter-max-bytes`, `stats-enabled.column.*`, keys added by future Iceberg versions), +metrics modes outside the supported set (`counts`, `none`, or unparseable values), any +`parquet.*` table property (including `parquet.enable.dictionary`), and any `parquet.*` key in +the session Hadoop configuration (with `HadoopFileIO`-backed output those reach iceberg-java's +writer but not the native one). Also gated explicitly: any `encryption.*` key, +`write.object-storage.enabled=true`, `write.location-provider.impl`, and `io-impl`. + +Other `write.*` properties are intentionally not gated because they cannot make the native +writer produce different data files: distribution and ordering settings shape the Spark plan +identically on both paths, WAP / branch / snapshot properties act on the JVM committer, +`write.avro.*` / `write.orc.*` apply only to formats already excluded, and merge-on-read +settings route the write through `WriteDelta`, which the split plan never intercepts. Every +rule is pinned by `CometIcebergWriteDetectionSuite`. + +Manifest `DataFile` metrics will be assembled on the JVM at commit time using Iceberg's own +`MetricsConfig` logic, so iceberg-java's metadata decisions — metrics modes, the +inferred-column cap (`write.metadata.metrics.max-inferred-column-defaults`), bound truncation, +and list/map bounds suppression — are respected exactly regardless of what the native writer +reports. The `counts`/`none` restrictions above remain only until that assembly lands. + +## Accepted divergences behind the toggle + +Some differences between parquet-mr and the pinned parquet-rs / iceberg-rust are unconditional — +they apply to every native write and cannot be configured away. Enabling +`nativeAcceleration.enabled` accepts them: + +- Footer key-value metadata differs: native files carry an `ARROW:schema` entry and no + `iceberg.schema` entry; iceberg-java files are the opposite. +- The Parquet root schema element is named `arrow_schema` (iceberg-java: `table`). +- `created_by` identifies parquet-rs, not parquet-mr. +- No page CRC checksums and no page-header statistics (parquet-mr writes both by default). +- Dictionary-encoded pages are labeled `RLE_DICTIONARY` (parquet-mr v1 files: `PLAIN_DICTIONARY`). +- Fixed-length binary columns (`uuid`, `fixed`, decimals with precision > 18) are not + dictionary-encoded (parquet-mr dictionary-encodes them). +- Row-group boundaries: parquet-mr flushes by byte size at a record-count check cadence, + parquet-rs buffers by row count. File rolling and file naming follow the same cadence-style + differences (iceberg-java checks the target file size every 1000 rows and names files + `---`; iceberg-rust checks per batch and uses a + process-local counter). +- Partition paths are not URL-escaped: iceberg-java percent-encodes partition directory names + and values (`region=a%2Fb`), iceberg-rust writes them raw (`region=a/b`). Readers resolve + files through manifest metadata, not paths, so query results are unaffected — but the + directory layout differs from iceberg-java's, and partition values containing characters + that are invalid in a URI (`:`, `#`, newline) may produce paths that `HadoopFileIO`-based + readers cannot open. +- Compressed page bytes are implementation-defined: the codec and any explicit level are + translated, but parquet-rs and parquet-mr embed different encoder implementations and + defaults (zstd default levels, LZ4 framing), so byte-identical output is not achievable even + for a default `zstd` table. The decompressed data is identical. For the same reason, + codec-level side channels (`zlib.compress.level`, `compression.brotli.quality`, + `io.compression.codec.zstd.level` — the last is present in every Hadoop configuration by + default) are not gated: they can only shift compressed bytes, which are already accepted as + divergent. + +All content not listed above — the logical data, encodings for non-FLBA columns, statistics +values, and manifest metadata — must match iceberg-java exactly, or the write falls back. diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index 6e7bb4271b..450afc430e 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -131,6 +131,16 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_ICEBERG_NATIVE_WRITE_ENABLED: ConfigEntry[Boolean] = + conf("spark.comet.write.iceberg.nativeAcceleration.enabled") + .category(CATEGORY_TESTING) + .doc( + "Whether to delegate the executor-side Parquet write to Comet's native (iceberg-rust) " + + "writer when the table's properties allow it. Requires " + + "`spark.comet.write.iceberg.splitOperator.enabled = true`. Off by default.") + .booleanConf + .createWithDefault(false) + val COMET_ICEBERG_DATA_FILE_CONCURRENCY_LIMIT: ConfigEntry[Int] = conf("spark.comet.scan.icebergNative.dataFileConcurrencyLimit") .category(CATEGORY_SCAN) diff --git a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala index edb321e8fc..078cb61f9d 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -58,6 +58,7 @@ object IcebergReflection extends Logging { val TABLE = "org.apache.iceberg.Table" val PARTITIONING = "org.apache.iceberg.Partitioning" val SPARK_WRITE = "org.apache.iceberg.spark.source.SparkWrite" + val TABLE_PROPERTIES = "org.apache.iceberg.TableProperties" // Iceberg 1.5.2 uses its own `ReplaceIcebergData` due to lack of `ReplaceData` in Spark 3.4. val REPLACE_ICEBERG_DATA = "org.apache.spark.sql.catalyst.plans.logical.ReplaceIcebergData" @@ -146,6 +147,29 @@ object IcebergReflection extends Logging { def isIcebergSparkWrite(write: Any): Boolean = sparkWriteClassOpt.exists(_.isInstance(write)) + def isIcebergBatchWrite(batchWrite: Any): Boolean = { + if (batchWrite == null) return false + batchWrite.getClass.getName.startsWith(ClassNames.SPARK_WRITE + "$") + } + + def getOuterSparkWrite(batchWrite: Any): Option[Any] = { + if (batchWrite == null) None + else { + try { + val field = batchWrite.getClass.getDeclaredField("this$0") + field.setAccessible(true) + Option(field.get(batchWrite)) + } catch { + case _: NoSuchFieldException => + None + case e: Exception => + logError( + s"Iceberg reflection failure: outer SparkWrite from BatchWrite: ${e.getMessage}") + None + } + } + } + def isReplaceIcebergData(plan: Any): Boolean = plan != null && plan.getClass.getName == ClassNames.REPLACE_ICEBERG_DATA @@ -1033,6 +1057,73 @@ object IcebergReflection extends Logging { getTableProperties(table).filter(_.containsKey("encryption.key-id")).map { props => Option(props.get("encryption.data-key-length")).map(_.toInt).getOrElse(16) } + + private def getSparkWriteField(sparkWrite: Any, fieldName: String): Option[Any] = + sparkWriteClassOpt.flatMap { cls => + try { + val field = cls.getDeclaredField(fieldName) + field.setAccessible(true) + Option(field.get(sparkWrite)) + } catch { + case _: NoSuchFieldException => None + case e: Exception => + logError( + s"Iceberg reflection failure: Failed to read SparkWrite.$fieldName: ${e.getMessage}") + None + } + } + + def getFormatFromSparkWrite(sparkWrite: Any): Option[String] = + getSparkWriteField(sparkWrite, "format") + .map(_.toString.toLowerCase(java.util.Locale.ROOT)) + + def getTableFromSparkWrite(sparkWrite: Any): Option[Any] = + getSparkWriteField(sparkWrite, "table") + + def getWritePropertiesFromSparkWrite(sparkWrite: Any): Option[Map[String, String]] = { + import scala.jdk.CollectionConverters._ + getSparkWriteField(sparkWrite, "writeProperties") + .map(_.asInstanceOf[java.util.Map[String, String]].asScala.toMap) + } + + private lazy val tablePropertiesClassOpt: Option[Class[_]] = + tryLoadClass(ClassNames.TABLE_PROPERTIES) + + def tablePropertyConstant(fieldName: String): String = + readTablePropertiesField(fieldName).asInstanceOf[String] + + def tablePropertyIntConstant(fieldName: String): Int = + readTablePropertiesField(fieldName).asInstanceOf[Integer].intValue() + + private def readTablePropertiesField(fieldName: String): Any = { + val cls = tablePropertiesClassOpt.getOrElse( + throw new IllegalStateException(s"${ClassNames.TABLE_PROPERTIES} is not on the classpath")) + try cls.getField(fieldName).get(null) + catch { + case e: NoSuchFieldException => + throw new IllegalStateException( + s"${ClassNames.TABLE_PROPERTIES}.$fieldName not found " + + "(unsupported Iceberg version?)", + e) + } + } + + def getDataLocation(table: Any): Option[String] = + try { + val locationProviderMethod = + findMethodInHierarchy(table.getClass, "locationProvider").getOrElse( + throw new NoSuchMethodException( + s"locationProvider() not found on ${table.getClass.getName}")) + val provider = locationProviderMethod.invoke(table) + val newDataLocMethod = provider.getClass.getMethod("newDataLocation", classOf[String]) + newDataLocMethod.setAccessible(true) + val location = newDataLocMethod.invoke(provider, "").asInstanceOf[String] + Some(location.stripSuffix("/")) + } catch { + case e: Exception => + logError(s"Iceberg reflection failure: Failed to get data location: ${e.getMessage}", e) + None + } } /** diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index ef2f37371c..cbf6c6a1e8 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -298,6 +298,9 @@ case class CometExecRule(session: SparkSession) case op: DataWritingCommandExec => convertToComet(op, CometDataWritingCommand).getOrElse(op) + case op: IcebergWriteExec if CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.get(op.conf) => + convertToComet(op, CometIcebergNativeWrite).getOrElse(op) + // For AQE broadcast stage on a Comet broadcast exchange case s @ BroadcastQueryStageExec(_, _: CometBroadcastExchangeExec, _) => convertToComet(s, CometExchangeSink).getOrElse(s) diff --git a/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala new file mode 100644 index 0000000000..5867bbb3f7 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala @@ -0,0 +1,325 @@ +/* + * 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.comet.serde.operator + +import java.util.Locale + +import scala.jdk.CollectionConverters._ +import scala.util.control.NonFatal + +import org.apache.hadoop.conf.Configuration +import org.apache.spark.internal.Logging +import org.apache.spark.sql.comet.{CometNativeExec, IcebergWriteExec} + +import org.apache.comet.{CometConf, ConfigEntry} +import org.apache.comet.CometSparkSessionExtensions.withFallbackReason +import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.serde.{CometOperatorSerde, Compatible, OperatorOuterClass, SupportLevel, Unsupported} +import org.apache.comet.serde.OperatorOuterClass.Operator + +object CometIcebergNativeWrite extends CometOperatorSerde[IcebergWriteExec] with Logging { + + override def enabledConfig: Option[ConfigEntry[Boolean]] = + Some(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED) + + override def requiresNativeChildren: Boolean = true + + object PropertyKeys { + lazy val ObjectStoreEnabled: String = + IcebergReflection.tablePropertyConstant("OBJECT_STORE_ENABLED") + lazy val WriteLocationProviderImpl: String = + IcebergReflection.tablePropertyConstant("WRITE_LOCATION_PROVIDER_IMPL") + lazy val DefaultWriteMetricsMode: String = + IcebergReflection.tablePropertyConstant("DEFAULT_WRITE_METRICS_MODE") + lazy val MetricsModeColumnPrefix: String = + IcebergReflection.tablePropertyConstant("METRICS_MODE_COLUMN_CONF_PREFIX") + lazy val BloomFilterColumnEnabledPrefix: String = + IcebergReflection.tablePropertyConstant("PARQUET_BLOOM_FILTER_COLUMN_ENABLED_PREFIX") + lazy val ParquetRowGroupCheckMinRecordCount: String = + IcebergReflection.tablePropertyConstant("PARQUET_ROW_GROUP_CHECK_MIN_RECORD_COUNT") + lazy val ParquetRowGroupCheckMinRecordCountDefault: Int = + IcebergReflection.tablePropertyIntConstant( + "PARQUET_ROW_GROUP_CHECK_MIN_RECORD_COUNT_DEFAULT") + lazy val ParquetRowGroupCheckMaxRecordCount: String = + IcebergReflection.tablePropertyConstant("PARQUET_ROW_GROUP_CHECK_MAX_RECORD_COUNT") + lazy val ParquetRowGroupCheckMaxRecordCountDefault: Int = + IcebergReflection.tablePropertyIntConstant( + "PARQUET_ROW_GROUP_CHECK_MAX_RECORD_COUNT_DEFAULT") + lazy val ParquetCompressionCodec: String = + IcebergReflection.tablePropertyConstant("PARQUET_COMPRESSION") + lazy val ParquetCompressionLevel: String = + IcebergReflection.tablePropertyConstant("PARQUET_COMPRESSION_LEVEL") + lazy val ParquetRowGroupSizeBytes: String = + IcebergReflection.tablePropertyConstant("PARQUET_ROW_GROUP_SIZE_BYTES") + lazy val ParquetPageSizeBytes: String = + IcebergReflection.tablePropertyConstant("PARQUET_PAGE_SIZE_BYTES") + lazy val ParquetPageRowLimit: String = + IcebergReflection.tablePropertyConstant("PARQUET_PAGE_ROW_LIMIT") + lazy val ParquetDictSizeBytes: String = + IcebergReflection.tablePropertyConstant("PARQUET_DICT_SIZE_BYTES") + val ParquetPageVersion: String = "write.parquet.page-version" + val ParquetPageVersionDefault: String = "v1" + val ParquetShredVariants: String = "write.parquet.shred-variants" + val ParquetVariantBufferSize: String = "write.parquet.variant-inference-buffer-size" + val ParquetEnableDictionary: String = "parquet.enable.dictionary" + val FileIOImpl: String = "io-impl" + } + + private val EncryptionPropertyPrefix = "encryption." + private val SupportedStorageSchemes: Set[String] = + Set("file", "memory", "s3", "s3a", "gs", "oss") + private val MinUnsupportedFormatVersion = 3 + private val ParquetWritePropertyPrefix = "write.parquet." + private val ParquetMrPropertyPrefix = "parquet." + + private lazy val vettedParquetWriteKeys: Set[String] = Set( + PropertyKeys.ParquetCompressionCodec, + PropertyKeys.ParquetCompressionLevel, + PropertyKeys.ParquetRowGroupSizeBytes, + PropertyKeys.ParquetPageSizeBytes, + PropertyKeys.ParquetPageRowLimit, + PropertyKeys.ParquetDictSizeBytes, + PropertyKeys.ParquetRowGroupCheckMinRecordCount, + PropertyKeys.ParquetRowGroupCheckMaxRecordCount, + PropertyKeys.ParquetPageVersion, + PropertyKeys.ParquetShredVariants, + PropertyKeys.ParquetVariantBufferSize) + + private lazy val vettedParquetWritePrefixes: Seq[String] = + Seq(PropertyKeys.BloomFilterColumnEnabledPrefix) + + override def getSupportLevel(op: IcebergWriteExec): SupportLevel = + try { + checkTriggers(op) match { + case Some(reason) => Unsupported(Some(reason)) + case None => Compatible(None) + } + } catch { + case NonFatal(e) => + Unsupported(Some(s"Iceberg native write detection failed: ${e.getMessage}")) + } + + private def checkTriggers(op: IcebergWriteExec): Option[String] = { + val batchWrite = op.batchWrite + if (!IcebergReflection.isIcebergBatchWrite(batchWrite)) { + return Some(s"not an Iceberg SparkWrite: ${batchWrite.getClass.getName}") + } + + val sparkWrite = IcebergReflection + .getOuterSparkWrite(batchWrite) + .getOrElse(return Some("could not unwrap SparkWrite")) + val table = IcebergReflection + .getTableFromSparkWrite(sparkWrite) + .getOrElse(return Some("SparkWrite.table is null")) + + val tableProperties = IcebergReflection + .getTableProperties(table) + .map(_.asScala.toMap) + .getOrElse(Map.empty[String, String]) + val writeProperties = IcebergReflection + .getWritePropertiesFromSparkWrite(sparkWrite) + .getOrElse(return Some("could not read SparkWrite.writeProperties")) + + val context = TriggerContext( + table, + tableProperties ++ writeProperties, + sparkWrite, + op.session.sessionState.newHadoopConf()) + triggers.iterator.map(rule => rule(context)).collectFirst { case Some(reason) => reason } + } + + private case class TriggerContext( + table: Any, + properties: Map[String, String], + sparkWrite: Any, + hadoopConf: Configuration) + + private type TriggerRule = TriggerContext => Option[String] + + private lazy val triggers: Seq[TriggerRule] = Seq( + requireFormatParquet, + requirePropertyAbsentOrNotTrue( + PropertyKeys.ObjectStoreEnabled, + "object-storage layout unsupported"), + requirePropertyAbsent( + PropertyKeys.WriteLocationProviderImpl, + "custom location provider unsupported"), + requireFormatVersionAtMostTwo, + requireNoEncryptionPrefix, + requireSupportedMetricsModes, + requireNoBloomFilterColumnsEnabled, + requireRowGroupCheckMinRecordCountAtDefault, + requireRowGroupCheckMaxRecordCountAtDefault, + requireParquetPageVersionDefault, + requireShredVariantsDisabled, + requireOnlyVettedParquetWriteProperties, + requirePropertyAbsent( + PropertyKeys.ParquetEnableDictionary, + "dictionary override unsupported"), + requireNoUnvettedParquetMrProperties, + requirePropertyAbsent(PropertyKeys.FileIOImpl, "custom FileIO unsupported"), + requireNoParquetHadoopConfOverrides, + requireSupportedStorageScheme) + + private val requireFormatParquet: TriggerRule = ctx => + IcebergReflection.getFormatFromSparkWrite(ctx.sparkWrite) match { + case None => Some("could not resolve the effective write format from SparkWrite") + case Some("parquet") => None + case Some(other) => Some(s"resolved write format=$other (only parquet is supported)") + } + + private def requirePropertyAbsentOrNotTrue(key: String, reason: String): TriggerRule = + ctx => { + if (ctx.properties.get(key).exists(_.equalsIgnoreCase("true"))) { + Some(s"$key=true ($reason)") + } else { + None + } + } + + private def requirePropertyAbsent(key: String, reason: String): TriggerRule = + ctx => { + if (ctx.properties.contains(key)) Some(s"$key is set ($reason)") else None + } + + private val requireFormatVersionAtMostTwo: TriggerRule = ctx => + IcebergReflection.getFormatVersion(ctx.table) match { + case Some(v) if v >= MinUnsupportedFormatVersion => Some(s"format-version=$v unsupported") + case Some(_) => None + case None => Some("could not determine the table format-version") + } + + private val requireNoEncryptionPrefix: TriggerRule = ctx => + ctx.properties.keys + .find(_.startsWith(EncryptionPropertyPrefix)) + .map(k => s"$k set: encryption unsupported") + + private val TruncateModePattern = """truncate\((\d+)\)""".r + + private def isSupportedMetricsMode(value: String): Boolean = + value.trim.toLowerCase(Locale.ROOT) match { + case "full" => true + case TruncateModePattern(n) => n.toInt > 0 + case _ => false + } + + private val requireSupportedMetricsModes: TriggerRule = ctx => { + val defaultKey = PropertyKeys.DefaultWriteMetricsMode + val prefix = PropertyKeys.MetricsModeColumnPrefix + ctx.properties + .find { case (k, v) => + (k == defaultKey || k.startsWith(prefix)) && !isSupportedMetricsMode(v) + } + .map { case (k, v) => s"$k=$v (supported metrics modes: full, truncate(N))" } + } + + private val requireNoBloomFilterColumnsEnabled: TriggerRule = ctx => { + val prefix = PropertyKeys.BloomFilterColumnEnabledPrefix + ctx.properties + .find { case (k, v) => k.startsWith(prefix) && v.equalsIgnoreCase("true") } + .map { case (k, _) => s"$k=true: bloom filters unsupported" } + } + + private val requireParquetPageVersionDefault: TriggerRule = ctx => { + val key = PropertyKeys.ParquetPageVersion + ctx.properties + .get(key) + .filter(_.trim.toLowerCase(Locale.ROOT) != PropertyKeys.ParquetPageVersionDefault) + .map(v => s"$key=$v unsupported") + } + + private val requireShredVariantsDisabled: TriggerRule = ctx => { + val key = PropertyKeys.ParquetShredVariants + ctx.properties + .get(key) + .filter(_.equalsIgnoreCase("true")) + .map(_ => s"$key=true (variant shredding changes the parquet schema)") + } + + private val requireOnlyVettedParquetWriteProperties: TriggerRule = ctx => + ctx.properties + .find { case (k, _) => + k.startsWith(ParquetWritePropertyPrefix) && + !vettedParquetWriteKeys.contains(k) && + !vettedParquetWritePrefixes.exists(k.startsWith) + } + .map { case (k, v) => s"$k=$v is not a vetted parquet write property" } + + private val requireNoUnvettedParquetMrProperties: TriggerRule = ctx => + ctx.properties.keys + .find(k => + k.startsWith(ParquetMrPropertyPrefix) && k != PropertyKeys.ParquetEnableDictionary) + .map(k => s"$k is set (parquet-mr properties are forwarded verbatim by iceberg-java)") + + private val requireNoParquetHadoopConfOverrides: TriggerRule = ctx => + ctx.hadoopConf.asScala + .map(_.getKey) + .find(_.startsWith(ParquetMrPropertyPrefix)) + .map(k => s"Hadoop configuration sets $k (reaches iceberg-java's writer but not native)") + + private val requireSupportedStorageScheme: TriggerRule = ctx => + IcebergReflection.getDataLocation(ctx.table) match { + case None => Some("could not resolve the table data location") + case Some(location) => + val scheme = if (location.contains("://")) { + location.substring(0, location.indexOf("://")).toLowerCase(Locale.ROOT) + } else { + "file" + } + if (SupportedStorageSchemes.contains(scheme)) None + else Some(s"unsupported storage scheme: $scheme") + } + + private lazy val requireRowGroupCheckMinRecordCountAtDefault: TriggerRule = + requireIntPropertyAtDefault( + PropertyKeys.ParquetRowGroupCheckMinRecordCount, + PropertyKeys.ParquetRowGroupCheckMinRecordCountDefault, + "row-group record-count cadence unsupported") + + private lazy val requireRowGroupCheckMaxRecordCountAtDefault: TriggerRule = + requireIntPropertyAtDefault( + PropertyKeys.ParquetRowGroupCheckMaxRecordCount, + PropertyKeys.ParquetRowGroupCheckMaxRecordCountDefault, + "row-group record-count cadence unsupported") + + private def requireIntPropertyAtDefault( + key: String, + default: Int, + reason: String): TriggerRule = ctx => + ctx.properties.get(key).flatMap { raw => + scala.util.Try(raw.trim.toInt).toOption match { + case Some(v) if v != default => Some(s"$key=$v (default=$default; $reason)") + case Some(_) => None + case None => Some(s"$key=$raw is not an int ($reason)") + } + } + + override def convert( + op: IcebergWriteExec, + builder: Operator.Builder, + childOp: Operator*): Option[OperatorOuterClass.Operator] = { + val _ = (builder, childOp) + withFallbackReason(op, "native Iceberg write is not yet implemented") + None + } + + override def createExec(nativeOp: Operator, op: IcebergWriteExec): CometNativeExec = + throw new UnsupportedOperationException("native Iceberg write not yet implemented") +} diff --git a/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala new file mode 100644 index 0000000000..d7a8d4dbdc --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteDetectionSuite.scala @@ -0,0 +1,685 @@ +/* + * 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.comet + +import java.io.File + +import org.apache.spark.SparkConf +import org.apache.spark.sql.CometTestBase +import org.apache.spark.sql.comet.IcebergWriteExec + +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.iceberg.IcebergReflection +import org.apache.comet.serde.{Compatible, SupportLevel, Unsupported} +import org.apache.comet.serde.operator.CometIcebergNativeWrite + +class CometIcebergWriteDetectionSuite extends CometTestBase with CometIcebergTestBase { + + override protected def sparkConf: SparkConf = { + super.sparkConf + .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true") + .set(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key, "true") + } + + test("clean parquet V2 table planned as AppendData yields Compatible") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "ok", partitionSpec = "") + assertSupportLevelIs[Compatible]("ok") + } + } + + test("registration tags a fall-back reason on the write exec") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "tagged", partitionSpec = "") + val writeExec = insertWriteExec("tagged") + val reasons = writeExec.getTagValue(CometExplainInfo.FALLBACK_REASONS) + assert( + reasons.exists(_.nonEmpty), + s"expected CometExecRule to record a fall-back reason on $writeExec") + } + } + + test("SparkWrite reflection helpers all resolve on the current Iceberg runtime") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "refl_probe", partitionSpec = "") + val sparkWrite = IcebergReflection + .getOuterSparkWrite(insertWriteExec("refl_probe").batchWrite) + .getOrElse(fail("could not unwrap outer SparkWrite from BatchWrite")) + val table = IcebergReflection + .getTableFromSparkWrite(sparkWrite) + .getOrElse(fail("SparkWrite.table reflection returned None")) + + assert(IcebergReflection.getFormatFromSparkWrite(sparkWrite).isDefined, "format") + assert( + IcebergReflection.getWritePropertiesFromSparkWrite(sparkWrite).isDefined, + "writeProperties") + assert(IcebergReflection.getDataLocation(table).isDefined, "dataLocation") + assert(IcebergReflection.getTableProperties(table).isDefined, "tableProperties") + } + } + + test("Compatible when a session conf overrides the compression codec") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "session_codec", partitionSpec = "") + withSQLConf("spark.sql.iceberg.compression-codec" -> "gzip") { + assertSupportLevelIs[Compatible]("session_codec") + } + } + } + + test("fall-back: write.format.default=orc") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "fmt_orc", + partitionSpec = "", + properties = Some("'write.format.default'='orc'")) + assertUnsupportedContains("fmt_orc", "format=orc", "only parquet") + } + } + + test("fall-back: per-write write-format option overrides parquet default") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "fmt_orc_opt", partitionSpec = "") + assertUnsupportedContains( + dfWriteExec("fmt_orc_opt", "write-format" -> "orc"), + "fmt_orc_opt", + "format=orc", + "only parquet") + } + } + + test("fall-back: write.object-storage.enabled=true") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "obj_store", + partitionSpec = "", + properties = Some("'write.object-storage.enabled'='true'")) + assertUnsupportedContains("obj_store", "write.object-storage.enabled") + } + } + + test("fall-back: write.location-provider.impl set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "loc_provider", + partitionSpec = "", + properties = Some("'write.location-provider.impl'='com.example.MyProvider'")) + assertUnsupportedContainsAllowingWriteFailure( + "loc_provider", + "write.location-provider.impl") + } + } + + test("fall-back: format-version=3") { + assume(icebergAvailable, "Iceberg not available in classpath") + assume(isSpark35Plus, "V3 tables require Iceberg 1.8.1+ (Spark 3.5 profile)") + withDetectionCatalog { dir => + createTable(dir, "v3", partitionSpec = "", properties = Some("'format-version'='3'")) + assertUnsupportedContains("v3", "format-version=3") + } + } + + test("fall-back: encryption.kms-client-impl set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "enc", + partitionSpec = "", + properties = Some("'encryption.kms-client-impl'='com.example.MyKms'")) + assertUnsupportedContainsAllowingWriteFailure("enc", "encryption") + } + } + + test("fall-back: write.metadata.metrics.default=counts") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "metrics_counts", + partitionSpec = "", + properties = Some("'write.metadata.metrics.default'='counts'")) + assertUnsupportedContains("metrics_counts", "write.metadata.metrics.default", "counts") + } + } + + test("fall-back: per-column metrics mode=counts") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "metrics_col_counts", + partitionSpec = "", + properties = Some("'write.metadata.metrics.column.id'='counts'")) + assertUnsupportedContains( + "metrics_col_counts", + "write.metadata.metrics.column.id", + "counts") + } + } + + test("fall-back: write.parquet.bloom-filter-max-bytes set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "bloom_max", + partitionSpec = "", + properties = Some("'write.parquet.bloom-filter-max-bytes'='524288'")) + assertUnsupportedContains("bloom_max", "write.parquet.bloom-filter-max-bytes") + } + } + + test("fall-back: per-column bloom filter enabled") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "bloom_col", + partitionSpec = "", + properties = Some("'write.parquet.bloom-filter-enabled.column.id'='true'")) + assertUnsupportedContains( + "bloom_col", + "write.parquet.bloom-filter-enabled.column.id", + "true") + } + } + + test("Compatible when the schema exceeds max-inferred-column-defaults") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "too_many_cols", + partitionSpec = "", + properties = Some("'write.metadata.metrics.max-inferred-column-defaults'='2'")) + assertSupportLevelIs[Compatible]("too_many_cols") + } + } + + test("fall-back: row-group-check-min-record-count non-default") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "rg_min", + partitionSpec = "", + properties = Some("'write.parquet.row-group-check-min-record-count'='500'")) + assertUnsupportedContains("rg_min", "write.parquet.row-group-check-min-record-count=500") + } + } + + test("Compatible when row-group-check-min-record-count is at default (100)") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "rg_min_default", + partitionSpec = "", + properties = Some("'write.parquet.row-group-check-min-record-count'='100'")) + assertSupportLevelIs[Compatible]("rg_min_default") + } + } + + test("fall-back: row-group-check-max-record-count non-default") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "rg_max", + partitionSpec = "", + properties = Some("'write.parquet.row-group-check-max-record-count'='50000'")) + assertUnsupportedContains("rg_max", "write.parquet.row-group-check-max-record-count=50000") + } + } + + test("fall-back: write.metadata.metrics.default=none") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "metrics_none", + partitionSpec = "", + properties = Some("'write.metadata.metrics.default'='none'")) + assertUnsupportedContains("metrics_none", "write.metadata.metrics.default", "none") + } + } + + test("fall-back: per-column metrics mode=none") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "col_metrics_none", + partitionSpec = "", + properties = Some("'write.metadata.metrics.column.region'='none'")) + assertUnsupportedContains( + "col_metrics_none", + "write.metadata.metrics.column.region", + "none") + } + } + + test("fall-back: write.parquet.page-version=v2") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "page_v2", + partitionSpec = "", + properties = Some("'write.parquet.page-version'='v2'")) + assertUnsupportedContains("page_v2", "write.parquet.page-version", "v2") + } + } + + test("fall-back: parquet.enable.dictionary set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "enable_dict", + partitionSpec = "", + properties = Some("'parquet.enable.dictionary'='false'")) + assertUnsupportedContains("enable_dict", "parquet.enable.dictionary") + } + } + + test("fall-back: per-column write.parquet.stats-enabled. set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "col_stats", + partitionSpec = "", + properties = Some("'write.parquet.stats-enabled.column.region'='false'")) + assertUnsupportedContains("col_stats", "write.parquet.stats-enabled.column.region") + } + } + + test("fall-back: unvetted write.parquet.* property") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "unvetted", + partitionSpec = "", + properties = Some("'write.parquet.bloom-filter-adaptive-enabled'='true'")) + assertUnsupportedContains( + "unvetted", + "write.parquet.bloom-filter-adaptive-enabled", + "not a vetted") + } + } + + test("fall-back: parquet.* table property other than enable.dictionary") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "pq_mr_prop", + partitionSpec = "", + properties = Some("'parquet.columnindex.truncate.length'='32'")) + assertUnsupportedContains("pq_mr_prop", "parquet.columnindex.truncate.length") + } + } + + test("Compatible when a codec level side-channel property is set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "codec_level", + partitionSpec = "", + properties = Some("'zlib.compress.level'='9'")) + assertSupportLevelIs[Compatible]("codec_level") + } + } + + test("fall-back: parquet.* key in the session Hadoop configuration") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "hadoop_conf", partitionSpec = "") + withSQLConf("parquet.block.size" -> "1048576") { + assertUnsupportedContains("hadoop_conf", "parquet.block.size", "Hadoop configuration") + } + } + } + + test("fall-back: io-impl set") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "io_impl", + partitionSpec = "", + properties = Some("'io-impl'='com.example.MyFileIO'")) + assertUnsupportedContainsAllowingWriteFailure("io_impl", "io-impl") + } + } + + test("fall-back: data location URI scheme not supported by the native writer") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "bad_scheme", + partitionSpec = "", + properties = Some("'write.data.path'='hdfs://nonexistent.invalid/iceberg/db/bad_scheme'")) + assertUnsupportedContainsAllowingWriteFailure("bad_scheme", "storage scheme", "hdfs") + } + } + + test("Compatible when the data location scheme is s3") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "s3_scheme", + partitionSpec = "", + properties = Some("'write.data.path'='s3://nonexistent-bucket/iceberg/db/s3_scheme'")) + assertSupportLevelIs[Compatible]("s3_scheme", allowWriteFailure = true) + } + } + + test("Compatible for the remaining supported data location schemes") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + Seq("gs", "oss", "memory").foreach { scheme => + val table = s"${scheme}_scheme" + createTable( + dir, + table, + partitionSpec = "", + properties = Some(s"'write.data.path'='$scheme://nonexistent/iceberg/db/$table'")) + assertSupportLevelIs[Compatible](table, allowWriteFailure = true) + } + } + } + + test("Compatible when the data location is an explicit file:// path") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "file_scheme", + partitionSpec = "", + properties = Some(s"'write.data.path'='file://${dir.getAbsolutePath}/file_scheme_data'")) + assertSupportLevelIs[Compatible]("file_scheme") + } + } + + test("fall-back: write.parquet.shred-variants=true") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "shred", + partitionSpec = "", + properties = Some("'write.parquet.shred-variants'='true'")) + assertUnsupportedContains("shred", "write.parquet.shred-variants") + } + } + + test("fall-back: unparseable write.metadata.metrics.default") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "metrics_typo", + partitionSpec = "", + properties = Some("'write.metadata.metrics.default'='truncat(16)'")) + assertUnsupportedContains("metrics_typo", "truncat(16)", "supported metrics modes") + } + } + + test("Compatible when write.spark.fanout.enabled=true") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "fanout", + partitionSpec = "PARTITIONED BY (bucket(4, id))", + properties = Some("'write.spark.fanout.enabled'='true'")) + assertSupportLevelIs[Compatible]("fanout") + } + } + + test("Compatible when write.target-file-size-bytes is non-default") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable( + dir, + "target_size", + partitionSpec = "", + properties = Some("'write.target-file-size-bytes'='1048576'")) + assertSupportLevelIs[Compatible]("target_size") + } + } + + test("no fall-back reason is recorded when nativeAcceleration is disabled") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "flag_off", partitionSpec = "") + withSQLConf(CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key -> "false") { + val writeExec = insertWriteExec("flag_off") + assert( + writeExec.getTagValue(CometExplainInfo.FALLBACK_REASONS).isEmpty, + "expected no fall-back reason on the write exec when the feature is disabled") + } + } + } + + test("fall-back: BatchWrite that is not an Iceberg SparkWrite") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "plain_write", partitionSpec = "") + val stub = new org.apache.spark.sql.connector.write.BatchWrite { + override def createBatchWriterFactory( + info: org.apache.spark.sql.connector.write.PhysicalWriteInfo) + : org.apache.spark.sql.connector.write.DataWriterFactory = + throw new UnsupportedOperationException("stub") + override def commit( + messages: Array[org.apache.spark.sql.connector.write.WriterCommitMessage]): Unit = + () + override def abort( + messages: Array[org.apache.spark.sql.connector.write.WriterCommitMessage]): Unit = + () + } + val fake = insertWriteExec("plain_write").copy(batchWrite = stub) + assertUnsupportedContains(fake, "plain_write", "not an Iceberg SparkWrite") + } + } + + test("Compatible when partitioned by a bucket transform") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "part_bucket", partitionSpec = "PARTITIONED BY (bucket(4, id))") + assertSupportLevelIs[Compatible]("part_bucket") + } + } + + test("Compatible when partitioned by identity on a string column") { + assume(icebergAvailable, "Iceberg not available in classpath") + withDetectionCatalog { dir => + createTable(dir, "part_string", partitionSpec = "PARTITIONED BY (region)") + assertSupportLevelIs[Compatible]("part_string") + } + } + + private val catalog = "cat" + private val ns = "db" + + private def withDetectionCatalog(f: File => Unit): Unit = withTempIcebergDir { warehouseDir => + withSQLConf( + s"spark.sql.catalog.$catalog" -> "org.apache.iceberg.spark.SparkCatalog", + s"spark.sql.catalog.$catalog.type" -> "hadoop", + s"spark.sql.catalog.$catalog.warehouse" -> warehouseDir.getAbsolutePath) { + f(warehouseDir) + } + } + + private def createTable( + warehouseDir: File, + tableName: String, + partitionSpec: String, + properties: Option[String] = None): Unit = { + val props = properties.map(s => s" TBLPROPERTIES ($s)").getOrElse("") + spark.sql(s""" + CREATE TABLE $catalog.$ns.$tableName ( + id INT, + region STRING, + amount DOUBLE + ) USING iceberg + $partitionSpec + $props + """) + } + + private def insertWriteExec( + tableName: String, + allowWriteFailure: Boolean = false): IcebergWriteExec = + captureWriteExec(tableName, allowWriteFailure) { + spark.sql(s"INSERT INTO $catalog.$ns.$tableName VALUES (1, 'us', 1.0)") + } + + private def dfWriteExec(tableName: String, options: (String, String)*): IcebergWriteExec = + captureWriteExec(tableName, allowWriteFailure = false) { + val df = spark + .createDataFrame(Seq((1, "us", 1.0))) + .toDF("id", "region", "amount") + val writer = options.foldLeft(df.writeTo(s"$catalog.$ns.$tableName")) { case (w, (k, v)) => + w.option(k, v) + } + writer.append() + } + + private def captureWriteExec(tableName: String, allowWriteFailure: Boolean)( + trigger: => Unit): IcebergWriteExec = { + val captured = + new java.util.concurrent.atomic.AtomicReference[org.apache.spark.sql.execution.SparkPlan]() + val listener = new org.apache.spark.sql.util.QueryExecutionListener { + override def onSuccess( + funcName: String, + qe: org.apache.spark.sql.execution.QueryExecution, + durationNs: Long): Unit = + captured.compareAndSet(null, qe.executedPlan) + override def onFailure( + funcName: String, + qe: org.apache.spark.sql.execution.QueryExecution, + exception: Exception): Unit = + captured.compareAndSet(null, qe.executedPlan) + } + var failure: Option[Throwable] = None + try org.apache.spark.CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + catch { case _: java.util.concurrent.TimeoutException => () } + spark.listenerManager.register(listener) + try { + try trigger + catch { case scala.util.control.NonFatal(t) => failure = Some(t) } + try org.apache.spark.CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + catch { case _: java.util.concurrent.TimeoutException => () } + } finally { + spark.listenerManager.unregister(listener) + } + if (!allowWriteFailure) { + failure.foreach(t => fail(s"write to $tableName failed unexpectedly", t)) + } + val plan = Option(captured.get()) + .getOrElse(fail(s"No QueryExecution captured for $tableName")) + findWriteExecOrFail(plan) + } + + private def findWriteExecOrFail( + plan: org.apache.spark.sql.execution.SparkPlan): IcebergWriteExec = + findWriteExec(plan).getOrElse(fail(s"no IcebergWriteExec found in:\n$plan")) + + private def findWriteExec( + plan: org.apache.spark.sql.execution.SparkPlan): Option[IcebergWriteExec] = + plan match { + case e: IcebergWriteExec => Some(e) + case other => + val descend = other.children.iterator ++ wrappedChildren(other).iterator + descend.flatMap(findWriteExec).toSeq.headOption + } + + private def wrappedChildren(plan: org.apache.spark.sql.execution.SparkPlan) + : Iterable[org.apache.spark.sql.execution.SparkPlan] = { + def viaAccessor(method: String): Option[org.apache.spark.sql.execution.SparkPlan] = + scala.util + .Try { + plan.getClass + .getMethod(method) + .invoke(plan) + .asInstanceOf[org.apache.spark.sql.execution.SparkPlan] + } + .toOption + .filter(_ ne plan) + Seq("commandPhysicalPlan", "executedPlan", "plan").flatMap(viaAccessor) + } + + private def assertSupportLevelIs[T <: SupportLevel: scala.reflect.ClassTag]( + tableName: String, + allowWriteFailure: Boolean = false): Unit = { + val support = + CometIcebergNativeWrite.getSupportLevel(insertWriteExec(tableName, allowWriteFailure)) + val expected = scala.reflect.classTag[T].runtimeClass + assert( + expected.isInstance(support), + s"expected ${expected.getSimpleName} for $tableName, got $support") + } + + private def assertUnsupportedContains(tableName: String, fragments: String*): Unit = + assertUnsupportedContains(insertWriteExec(tableName), tableName, fragments: _*) + + private def assertUnsupportedContainsAllowingWriteFailure( + tableName: String, + fragments: String*): Unit = + assertUnsupportedContains( + insertWriteExec(tableName, allowWriteFailure = true), + tableName, + fragments: _*) + + private def assertUnsupportedContains( + writeExec: IcebergWriteExec, + tableName: String, + fragments: String*): Unit = { + val support = CometIcebergNativeWrite.getSupportLevel(writeExec) + support match { + case Unsupported(Some(reason)) => + fragments.foreach(f => + assert(reason.contains(f), s"reason '$reason' missing fragment '$f' for $tableName")) + case Unsupported(None) => + fail(s"Unsupported without a reason string for $tableName") + case other => + fail(s"expected Unsupported for $tableName, got $other") + } + } +}