jordepic commented on code in PR #5361:
URL: https://github.com/apache/datafusion-comet/pull/5361#discussion_r3914714424
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala:
##########
@@ -172,21 +169,23 @@ object CometIcebergNativeWrite extends
CometOperatorSerde[IcebergWriteExec] with
PropertyKeys.WriteLocationProviderImpl,
"custom location provider unsupported"),
requireFormatVersionAtMostTwo,
+ requireNoUuidColumns,
requireNoEncryptionPrefix,
- requireSupportedMetricsModes,
requireNoBloomFilterColumnsEnabled,
requireRowGroupCheckMinRecordCountAtDefault,
requireRowGroupCheckMaxRecordCountAtDefault,
requireParquetPageVersionDefault,
requireShredVariantsDisabled,
+ requireParseableCompressionLevel,
requireOnlyVettedParquetWriteProperties,
requirePropertyAbsent(
PropertyKeys.ParquetEnableDictionary,
"dictionary override unsupported"),
requireNoUnvettedParquetMrProperties,
requirePropertyAbsent(PropertyKeys.FileIOImpl, "custom FileIO
unsupported"),
Review Comment:
Good catch — the property rule genuinely cannot see a catalog-level
`io-impl` (or a custom catalog installing its own FileIO). Added
`requireRecognizedTableFileIO`: the gate now inspects the instantiated
`table.io()` against the same class-hierarchy allowlist the scan side uses
(`COMPATIBLE_FILE_IO_CLASSES`), with one deliberate difference — the
`EncryptingFileIO` family, which the scan accepts, is rejected on the write
side because the native writer produces plaintext data files. The property rule
stays as a cheap early signal. New detection test registers a catalog with a
catalog-level `io-impl` pointing at a FileIO that delegates to `HadoopFileIO`
by composition (a subclass would pass the hierarchy check by design) and pins
the fallback reason.
##########
spark/src/main/scala/org/apache/comet/serde/operator/IcebergWriteProtoTranslation.scala:
##########
@@ -0,0 +1,193 @@
+/*
+ * 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.Try
+
+import org.apache.comet.iceberg.IcebergReflection
+import org.apache.comet.serde.OperatorOuterClass._
+
+/**
+ * Pure translation from the resolved per-write property map (Iceberg table
properties merged with
+ * any write options) and a small set of driver-side state values into the
protobuf messages
+ * shipped to the Rust writer.
+ *
+ * Kept free of any `SparkWrite` reference so each translation function can be
unit-tested with
+ * plain `Map[String, String]` inputs. The serde wires these into the protobuf
during conversion;
+ * the JVM exec wrapper fills in `partition_id` and `task_attempt_id` at task
launch.
+ */
+object IcebergWriteProtoTranslation {
+
+ /**
+ * Iceberg `TableProperties` constants the translation depends on. Resolved
lazily through the
+ * reflection bridge so we always quote Iceberg's canonical names rather
than duplicating
+ * literal strings.
+ */
+ object Keys {
+ lazy val ParquetCompression: String =
+ IcebergReflection.tablePropertyConstant("PARQUET_COMPRESSION")
+ lazy val ParquetCompressionDefaultSince14: String =
+
IcebergReflection.tablePropertyConstant("PARQUET_COMPRESSION_DEFAULT_SINCE_1_4_0")
+ 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")
+ }
+
+ /** Iceberg's numeric defaults, pulled at runtime so they stay in lock-step
with the runtime. */
+ object Defaults {
+ lazy val RowGroupSizeBytes: Long =
+
IcebergReflection.tablePropertyIntConstant("PARQUET_ROW_GROUP_SIZE_BYTES_DEFAULT").toLong
+ lazy val PageSizeBytes: Long =
+
IcebergReflection.tablePropertyIntConstant("PARQUET_PAGE_SIZE_BYTES_DEFAULT").toLong
+ lazy val DictSizeBytes: Long =
+
IcebergReflection.tablePropertyIntConstant("PARQUET_DICT_SIZE_BYTES_DEFAULT").toLong
+ lazy val PageRowLimit: Int =
+
IcebergReflection.tablePropertyIntConstant("PARQUET_PAGE_ROW_LIMIT_DEFAULT")
+ }
+
+ /** Builds the parquet settings message. Pure: no SparkWrite or Iceberg
`Table` access. */
+ def buildParquetSettings(
+ props: Map[String, String],
+ createdBy: String): IcebergParquetWriteSettings = {
+ val rowGroupSize =
+ parseLong(props, Keys.ParquetRowGroupSizeBytes,
Defaults.RowGroupSizeBytes)
+ val pageSize = parseLong(props, Keys.ParquetPageSizeBytes,
Defaults.PageSizeBytes)
+ val dictSize = parseLong(props, Keys.ParquetDictSizeBytes,
Defaults.DictSizeBytes)
+ val pageRowLimit = parseInt(props, Keys.ParquetPageRowLimit,
Defaults.PageRowLimit)
+ val compression = resolveCompression(props)
+ val builder = IcebergParquetWriteSettings
+ .newBuilder()
+ .setCompression(compression)
+ .setRowGroupSizeBytes(rowGroupSize)
+ .setPageSizeBytes(pageSize)
+ .setDictSizeBytes(dictSize)
+ .setPageRowLimit(pageRowLimit)
+ .setCreatedBy(createdBy)
+
+ resolveCompressionLevel(props,
compression).foreach(builder.setCompressionLevel)
+
+ builder.build()
+ }
+
+ /**
+ * Driver-side resolution of the iceberg-rust writer flavor. The choice
mirrors
+ * `SparkWrite$WriterFactory`: unpartitioned tables get
`UnpartitionedDataWriter`; partitioned
+ * tables either fan out or cluster based on
`SparkWriteConf.useFanoutWriter`.
+ */
+ def resolveWriterMode(
+ specIsUnpartitioned: Boolean,
+ useFanoutWriter: Boolean): IcebergWriterMode = {
+ if (specIsUnpartitioned) IcebergWriterMode.ICEBERG_WRITER_UNPARTITIONED
+ else if (useFanoutWriter) IcebergWriterMode.ICEBERG_WRITER_FANOUT
+ else IcebergWriterMode.ICEBERG_WRITER_CLUSTERED
+ }
+
+ /** Builds the per-write broadcast message. */
+ // scalastyle:off argcount
+ def buildCommon(
+ catalogProperties: Map[String, String],
+ metadataLocation: String,
+ icebergSchemaJson: String,
+ partitionSpecJson: String,
+ sortOrderId: Int,
+ dataLocation: String,
+ operationId: String,
+ targetFileSizeBytes: Long,
+ writerMode: IcebergWriterMode,
+ parquetSettings: IcebergParquetWriteSettings,
+ catalogName: Option[String]): IcebergWriteCommon = {
+ val builder = IcebergWriteCommon
+ .newBuilder()
+ .setMetadataLocation(metadataLocation)
+ .setIcebergSchemaJson(icebergSchemaJson)
+ .setPartitionSpecJson(partitionSpecJson)
+ .setSortOrderId(sortOrderId)
+ .setDataLocation(dataLocation)
+ .setOperationId(operationId)
+ .setTargetFileSizeBytes(targetFileSizeBytes)
+ .setWriterMode(writerMode)
+ .setParquetSettings(parquetSettings)
+ if (catalogProperties.nonEmpty)
builder.putAllCatalogProperties(catalogProperties.asJava)
+ catalogName.foreach(builder.setCatalogName)
+ builder.build()
+ }
+ // scalastyle:on argcount
+
+ // --- Internals ------------------------------------------------------------
+
+ /**
+ * Iceberg accepts `uncompressed`, `none` (treated as uncompressed),
`snappy`, `gzip`, `lz4`,
+ * `zstd`, `brotli`. Defaults to `zstd` (since Iceberg 1.4).
+ */
+ private def resolveCompression(props: Map[String, String]): CompressionCodec
= {
+ val raw = props
+ .get(Keys.ParquetCompression)
+ .map(_.trim.toLowerCase(Locale.ROOT))
+ .getOrElse(Keys.ParquetCompressionDefaultSince14)
+ raw match {
+ case "uncompressed" | "none" => CompressionCodec.None
+ case "snappy" => CompressionCodec.Snappy
+ case "gzip" => CompressionCodec.Gzip
+ case "lz4" => CompressionCodec.Lz4
+ case "zstd" => CompressionCodec.Zstd
+ case "brotli" => CompressionCodec.Brotli
+ case other =>
+ throw new IllegalArgumentException(s"Unsupported parquet codec
'$other'")
+ }
+ }
+
+ /**
+ * Iceberg leaves `write.parquet.compression-level` null by default and lets
each parquet writer
+ * pick its own per-codec default. The one known divergence between
parquet-rs and parquet-mr is
+ * zstd (parquet-rs default 1, parquet-mr default 3). To produce files the
same size as
+ * iceberg-java would, substitute parquet-mr's 3 when zstd is in use and the
user did not set an
+ * explicit level. gzip defaults to 6 on both sides; snappy and lz4 have no
level concept;
+ * compressed bytes are an accepted divergence regardless (see
iceberg-writes.md).
+ */
+ private def resolveCompressionLevel(
+ props: Map[String, String],
+ codec: CompressionCodec): Option[Int] = {
+ val explicit =
+ props.get(Keys.ParquetCompressionLevel).flatMap(s =>
Try(s.trim.toInt).toOption)
+ explicit.orElse(parquetMrDefaultLevel(codec))
+ }
+
+ private def parquetMrDefaultLevel(codec: CompressionCodec): Option[Int] =
codec match {
+ case CompressionCodec.Zstd => Some(3)
+ case _ => None
+ }
+
+ private def parseLong(props: Map[String, String], key: String, default:
Long): Long =
+ props.get(key).flatMap(s => Try(s.trim.toLong).toOption).getOrElse(default)
Review Comment:
Right on all counts — iceberg-java reads these through
`PropertyUtil.propertyAsInt` (`Integer.parseInt` semantics: no trimming,
overflow throws) and parquet-mr rejects non-positive values at write time,
while the native translation was silently substituting defaults. Fixed in both
layers: a new gate rule (`requirePositiveIntParquetSizes`) declines any of the
four size/limit properties that is not a positive Java int, so those values
fail on the stock path exactly as they would without Comet; and the translation
helpers now use strict `Integer.parseInt` with no silent fallback (a parse
failure there would be a gate/translation mismatch, and throws loudly). The
compression-level rule was aligned to the same no-trim semantics. Boundary
tests cover `garbage`, `0`, `-1`, `2147483648`, and `" 1024"` per key in the
detection suite, plus strict-parse pins in `IcebergWriteProtoTranslationSuite`.
##########
native/core/src/execution/operators/iceberg_common.rs:
##########
@@ -0,0 +1,142 @@
+// 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.
+
+//! Helpers shared between the Iceberg scan and Iceberg write operators.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use datafusion::common::DataFusionError;
+use iceberg::io::{FileIO, FileIOBuilder, StorageFactory};
+use iceberg_storage_opendal::{CustomAwsCredentialLoader,
OpenDalStorageFactory};
+
+use crate::cloud::s3::credential_bridge::{AccessMode, CometS3CredentialBridge};
+
+/// Activation key for the `CometS3CredentialProvider` SPI, read from a
catalog's `s3.*` property
+/// bag.
+const ICEBERG_PROVIDER_CLASS_PROPERTY: &str =
"s3.comet.credential.provider.class";
+
+/// Key prefixes forwarded to iceberg-rust's `FileIO`. The full unfiltered
catalog bag (catalog
+/// URI, OAuth tokens, credentials.uri, tenant-id, etc.) is kept upstream so
+/// `CometS3CredentialBridge` can read whatever the vendor needs.
+const STORAGE_PROPERTY_PREFIXES: &[&str] = &["s3.", "gcs.", "adls.",
"client."];
Review Comment:
You are right that this was aspirational: `oss.*` was never in the forwarded
property prefixes and nothing tests the path. Took your second option — `oss`
is removed from the supported schemes (gate, docs, and the rust factory arm,
which nothing else could reach since the scan side never admits an OSS FileIO),
with a detection test pinning the fallback. Forwarding `oss.*` plus a
functional test can bring it back as a follow-up if someone has an OSS
environment to validate against.
##########
spark/src/main/scala/org/apache/comet/serde/operator/CometIcebergNativeWrite.scala:
##########
@@ -216,29 +215,35 @@ object CometIcebergNativeWrite extends
CometOperatorSerde[IcebergWriteExec] with
case None => Some("could not determine the table format-version")
}
+ // Iceberg maps `uuid` to Spark's StringType, so the native writer would
receive a Utf8 column
+ // while iceberg-rust's target Arrow schema demands FixedSizeBinary(16) --
no Arrow cast bridges
+ // the two, so the write would pass detection and then fail the task.
Decline it up front. This
+ // is the only Spark-writable Iceberg type with such a mismatch: `fixed(N)`
arrives as Binary
+ // and casts to FixedSizeBinary(N), and the V3-only types are excluded by
the format-version
+ // gate.
+ private val requireNoUuidColumns: TriggerRule = ctx =>
+ IcebergReflection
+ .getWriteSchemaFromSparkWrite(ctx.sparkWrite)
+ .orElse(IcebergReflection.getSchema(ctx.table)) match {
+ case None => Some("could not resolve the write schema for column type
checking")
+ case Some(schema) =>
+ IcebergReflection
+ .findFieldWithTypeIds(schema, UnsupportedWriteTypeIds)
+ .map { case (name, typeId) =>
+ s"column $name has Iceberg type
${typeId.toLowerCase(Locale.ROOT)}, " +
+ "which the native writer cannot reproduce"
+ }
+ }
+
private val requireNoEncryptionPrefix: TriggerRule = ctx =>
ctx.properties.keys
.find(_.startsWith(EncryptionPropertyPrefix))
Review Comment:
Agreed — the property prefix is a proxy, and a custom `TableOperations` can
install an `EncryptionManager` without any table property. Added
`requirePlaintextEncryptionManager`: the gate now reflects `table.encryption()`
and declines anything whose class is not
`org.apache.iceberg.encryption.PlaintextEncryptionManager`, failing closed when
the accessor cannot be resolved at all. The `encryption.*` property rule stays
as an early signal. (Same-commit related change: the new `table.io()` check
also rejects the `EncryptingFileIO` family on the write side, so encrypted
tables are declined from both directions.) Unit tests in
`IcebergReflectionSuite` cover the plaintext, custom-manager, and unresolvable
cases.
##########
native/core/src/execution/operators/iceberg_common.rs:
##########
@@ -0,0 +1,142 @@
+// 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.
+
+//! Helpers shared between the Iceberg scan and Iceberg write operators.
+
+use std::collections::HashMap;
+use std::sync::Arc;
+
+use datafusion::common::DataFusionError;
+use iceberg::io::{FileIO, FileIOBuilder, StorageFactory};
+use iceberg_storage_opendal::{CustomAwsCredentialLoader,
OpenDalStorageFactory};
+
+use crate::cloud::s3::credential_bridge::{AccessMode, CometS3CredentialBridge};
+
+/// Activation key for the `CometS3CredentialProvider` SPI, read from a
catalog's `s3.*` property
+/// bag.
+const ICEBERG_PROVIDER_CLASS_PROPERTY: &str =
"s3.comet.credential.provider.class";
+
+/// Key prefixes forwarded to iceberg-rust's `FileIO`. The full unfiltered
catalog bag (catalog
+/// URI, OAuth tokens, credentials.uri, tenant-id, etc.) is kept upstream so
+/// `CometS3CredentialBridge` can read whatever the vendor needs.
+const STORAGE_PROPERTY_PREFIXES: &[&str] = &["s3.", "gcs.", "adls.",
"client."];
+
+/// Pick an OpenDAL storage backend from a URI's scheme. `file` (or no scheme)
falls through to
+/// the local file system. `memory` is used by the write path to assemble
manifest bytes that
+/// stay entirely in-process. For S3, the Comet credential bridge is wired in
when a provider
+/// class is configured; `access_mode` is forwarded to the JVM SPI so the read
and write paths can
+/// be granted different (e.g. read-only vs read-write) credentials.
+pub(crate) fn storage_factory_for(
+ path: &str,
+ catalog_properties: &HashMap<String, String>,
+ catalog_name: &str,
+ access_mode: AccessMode,
+) -> Result<Arc<dyn StorageFactory>, DataFusionError> {
+ let scheme = if path.contains("://") {
+ path.split("://").next().unwrap_or("file")
+ } else {
+ "file"
+ };
+ match scheme {
+ "file" => Ok(Arc::new(OpenDalStorageFactory::Fs)),
+ "memory" => Ok(Arc::new(OpenDalStorageFactory::Memory)),
+ "s3" | "s3a" => {
+ let customized_credential_load =
+ build_s3_credential_loader(path, catalog_properties,
catalog_name, access_mode);
+ Ok(Arc::new(OpenDalStorageFactory::S3 {
+ customized_credential_load,
+ }))
+ }
+ "gs" => Ok(Arc::new(OpenDalStorageFactory::Gcs)),
+ "oss" => Ok(Arc::new(OpenDalStorageFactory::Oss)),
+ _ => Err(DataFusionError::Execution(format!(
+ "Unsupported storage scheme: {scheme}"
+ ))),
+ }
+}
+
+/// Build a `FileIO` whose storage scheme is inferred from `reference_path`
and whose properties
+/// come from the catalog. The reference path is the metadata location for
reads or the data
+/// location for writes — anything that carries the right URI scheme.
`catalog_name` is the
+/// credential dispatch key and `access_mode` is the access intent forwarded
to the S3 credential
+/// bridge, so the write path can request write-capable credentials.
+pub(crate) fn load_file_io(
+ catalog_properties: &HashMap<String, String>,
+ reference_path: &str,
+ catalog_name: &str,
+ access_mode: AccessMode,
+) -> Result<FileIO, DataFusionError> {
+ let factory = storage_factory_for(
+ reference_path,
+ catalog_properties,
+ catalog_name,
+ access_mode,
+ )?;
+ let mut file_io_builder = FileIOBuilder::new(factory);
+
+ // Narrow to storage-prefix keys before forwarding to iceberg-rust's
FileIO. The full
+ // unfiltered bag (catalog URI, OAuth tokens, credentials.uri, tenant-id,
etc.) is kept
+ // upstream so CometS3CredentialBridge can read whatever the vendor needs.
+ for (key, value) in catalog_properties {
+ if STORAGE_PROPERTY_PREFIXES.iter().any(|p| key.starts_with(p)) {
+ file_io_builder = file_io_builder.with_prop(key, value);
+ }
+ }
+
+ Ok(file_io_builder.build())
+}
+
+/// Wires the configured Comet credential provider into opendal's S3 service,
or returns `None`
+/// so opendal falls back to its default credential chain.
+fn build_s3_credential_loader(
+ reference_path: &str,
+ catalog_properties: &HashMap<String, String>,
+ catalog_name: &str,
+ access_mode: AccessMode,
+) -> Option<CustomAwsCredentialLoader> {
+ let url = url::Url::parse(reference_path).ok()?;
+ let bucket = url.host_str()?;
+ let provider_class = catalog_properties
+ .get(ICEBERG_PROVIDER_CLASS_PROPERTY)
+ .map(|s| s.trim())
+ .filter(|s| !s.is_empty())?;
+ // Fall back to the bucket when the table has no catalog identity (e.g.
HadoopTables loaded by
+ // raw path).
+ let dispatch_key: &str = if catalog_name.is_empty() {
+ bucket
+ } else {
+ catalog_name
+ };
+ let bridge = CometS3CredentialBridge::new(
+ provider_class,
+ dispatch_key,
+ bucket,
+ url.path(),
+ access_mode,
+ catalog_properties,
+ );
+ match bridge {
+ Ok(b) => Some(CustomAwsCredentialLoader::new(b)),
+ Err(e) => {
+ log::warn!(
+ "Failed to initialize CometS3CredentialBridge for
{provider_class}: {e}; \
+ falling back to default opendal credential chain"
Review Comment:
Agreed — silently switching which credentials perform a write after the
configured provider failed is not acceptable. `build_s3_credential_loader` now
returns `Result<Option<...>>`: with a provider explicitly configured,
`AccessMode::Write` fails closed with an error naming the provider class, while
`AccessMode::Read` keeps the warn-and-fall-back behavior (a wrong-credential
read fails on permissions rather than corrupting anything, and changing scan
behavior felt out of scope here). No provider configured still means `Ok(None)`
→ default chain, unchanged.
--
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]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
