fix: Fall back to native_comet when object store not supported by native_iceberg_compat (#2251)

Author: andygrove

common/src/main/java/org/apache/comet/parquet/Native.java

@@ -241,6 +241,13 @@ public static native void setPageV2(
   // to arrow.
   //      Add batch size, datetimeRebaseModeSpec, metrics(how?)...
 
+  /**
+   * Verify that object store options are valid. An exception will be thrown if the provided options
+   * are not valid.
+   */
+  public static native void validateObjectStoreConfig(
+      String filePath, Map<String, String> objectStoreOptions);
+
   /**
    * Initialize a record batch reader for a PartitionedFile
    *

docs/source/user-guide/compatibility.md

@@ -73,6 +73,11 @@ The `native_datafusion` scan has some additional limitations:
 [#1545]: https://github.com/apache/datafusion-comet/issues/1545
 [#1758]: https://github.com/apache/datafusion-comet/issues/1758
 
+### S3 Support with `native_iceberg_compat`
+
+- When using the default AWS S3 endpoint (no custom endpoint configured), a valid region is required. Comet 
+  will attempt to resolve the region if it is not provided.
+
 ## ANSI Mode
 
 Comet will fall back to Spark for the following expressions when ANSI mode is enabled, unless

native/core/src/parquet/mod.rs

@@ -666,6 +666,35 @@ pub fn get_object_store_options(
     Ok(collected_map)
 }
 
+/// # Safety
+/// This function is inherently unsafe since it deals with raw pointers passed from JNI.
+#[no_mangle]
+pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_validateObjectStoreConfig(
+    e: JNIEnv,
+    _jclass: JClass,
+    file_path: jstring,
+    object_store_options: jobject,
+) {
+    try_unwrap_or_throw(&e, |mut env| unsafe {
+        let session_config = SessionConfig::new();
+        let planner =
+            PhysicalPlanner::new(Arc::new(SessionContext::new_with_config(session_config)), 0);
+        let session_ctx = planner.session_ctx();
+        let path: String = env
+            .get_string(&JString::from_raw(file_path))
+            .unwrap()
+            .into();
+        let object_store_config =
+            get_object_store_options(&mut env, JObject::from_raw(object_store_options))?;
+        let (_, _) = prepare_object_store_with_configs(
+            session_ctx.runtime_env(),
+            path.clone(),
+            &object_store_config,
+        )?;
+        Ok(())
+    })
+}
+
 /// # Safety
 /// This function is inherently unsafe since it deals with raw pointers passed from JNI.
 #[no_mangle]
@@ -686,9 +715,9 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
 ) -> jlong {
     try_unwrap_or_throw(&e, |mut env| unsafe {
         let session_config = SessionConfig::new().with_batch_size(batch_size as usize);
-        let planer =
+        let planner =
             PhysicalPlanner::new(Arc::new(SessionContext::new_with_config(session_config)), 0);
-        let session_ctx = planer.session_ctx();
+        let session_ctx = planner.session_ctx();
 
         let path: String = env
             .get_string(&JString::from_raw(file_path))
@@ -716,7 +745,7 @@ pub unsafe extern "system" fn Java_org_apache_comet_parquet_Native_initRecordBat
             let filter_buffer = env.convert_byte_array(&filter_array)?;
             let filter_expr = serde::deserialize_expr(filter_buffer.as_slice())?;
             Some(vec![
-                planer.create_expr(&filter_expr, Arc::clone(&data_schema))?
+                planner.create_expr(&filter_expr, Arc::clone(&data_schema))?
             ])
         } else {
             None

native/core/src/parquet/objectstore/s3.rs

@@ -68,6 +68,12 @@ pub fn create_store(
     }
     let path = Path::parse(path)?;
 
+    if configs.contains_key("fs.s3a.endpoint") {
+        return Err(object_store::Error::NotSupported {
+            source: "Custom S3 endpoints are not supported".into(),
+        });
+    }
+
     let mut builder = AmazonS3Builder::new()
         .with_url(url.to_string())
         .with_allow_http(true);

spark/src/main/scala/org/apache/comet/rules/CometScanRule.scala

@@ -19,8 +19,13 @@
 
 package org.apache.comet.rules
 
+import java.net.URI
+
+import scala.collection.{mutable, JavaConverters}
 import scala.collection.mutable.ListBuffer
 
+import org.apache.hadoop.conf.Configuration
+import org.apache.spark.internal.Logging
 import org.apache.spark.sql.SparkSession
 import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, GenericInternalRow, PlanExpression}
 import org.apache.spark.sql.catalyst.rules.Rule
@@ -34,17 +39,19 @@ import org.apache.spark.sql.execution.datasources.v2.parquet.ParquetScan
 import org.apache.spark.sql.internal.SQLConf
 import org.apache.spark.sql.types._
 
-import org.apache.comet.{CometConf, DataTypeSupport}
+import org.apache.comet.{CometConf, CometNativeException, DataTypeSupport}
 import org.apache.comet.CometConf._
 import org.apache.comet.CometSparkSessionExtensions.{isCometLoaded, isCometScanEnabled, withInfo, withInfos}
 import org.apache.comet.DataTypeSupport.isComplexType
-import org.apache.comet.parquet.{CometParquetScan, SupportsComet}
+import org.apache.comet.objectstore.NativeConfig
+import org.apache.comet.parquet.{CometParquetScan, Native, SupportsComet}
 import org.apache.comet.shims.CometTypeShim
 
 /**
  * Spark physical optimizer rule for replacing Spark scans with Comet scans.
  */
 case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with CometTypeShim {
+  import CometScanRule._
 
   private lazy val showTransformations = CometConf.COMET_EXPLAIN_TRANSFORMATIONS.get()
 
@@ -295,8 +302,17 @@ case class CometScanRule(session: SparkSession) extends Rule[SparkPlan] with Com
     val fallbackReasons = new ListBuffer[String]()
 
     // native_iceberg_compat only supports local filesystem and S3
-    if (!scanExec.relation.inputFiles
+    if (scanExec.relation.inputFiles
         .forall(path => path.startsWith("file://") || path.startsWith("s3a://"))) {
+
+      val filePath = scanExec.relation.inputFiles.headOption
+      if (filePath.exists(_.startsWith("s3a://"))) {
+        validateObjectStoreConfig(
+          filePath.get,
+          session.sparkContext.hadoopConfiguration,
+          fallbackReasons)
+      }
+    } else {
       fallbackReasons += s"$SCAN_NATIVE_ICEBERG_COMPAT only supports local filesystem and S3"
     }
 
@@ -376,3 +392,64 @@ case class CometScanTypeChecker(scanImpl: String) extends DataTypeSupport with C
     }
   }
 }
+
+object CometScanRule extends Logging {
+
+  /**
+   * Validating object store configs can cause requests to be made to S3 APIs (such as when
+   * resolving the region for a bucket). We use a cache to reduce the number of S3 calls.
+   *
+   * The key is the config map converted to a string. The value is the reason that the config is
+   * not valid, or None if the config is valid.
+   */
+  val configValidityMap = new mutable.HashMap[String, Option[String]]()
+
+  /**
+   * We do not expect to see a large number of unique configs within the lifetime of a Spark
+   * session, but we reset the cache once it reaches a fixed size to prevent it growing
+   * indefinitely.
+   */
+  val configValidityMapMaxSize = 1024
+
+  def validateObjectStoreConfig(
+      filePath: String,
+      hadoopConf: Configuration,
+      fallbackReasons: mutable.ListBuffer[String]): Unit = {
+    val objectStoreConfigMap =
+      NativeConfig.extractObjectStoreOptions(hadoopConf, URI.create(filePath))
+
+    val cacheKey = objectStoreConfigMap
+      .map { case (k, v) =>
+        s"$k=$v"
+      }
+      .toList
+      .sorted
+      .mkString("\n")
+
+    if (configValidityMap.size >= configValidityMapMaxSize) {
+      logWarning("Resetting S3 object store validity cache")
+      configValidityMap.clear()
+    }
+
+    configValidityMap.get(cacheKey) match {
+      case Some(Some(reason)) =>
+        fallbackReasons += reason
+        throw new CometNativeException(reason)
+      case Some(None) =>
+      // previously validated
+      case _ =>
+        try {
+          val objectStoreOptions = JavaConverters.mapAsJavaMap(objectStoreConfigMap)
+          Native.validateObjectStoreConfig(filePath, objectStoreOptions)
+        } catch {
+          case e: Exception =>
+            val reason = "Object store config not supported by " +
+              s"$SCAN_NATIVE_ICEBERG_COMPAT: ${e.getMessage}"
+            fallbackReasons += reason
+            configValidityMap.put(cacheKey, Some(reason))
+            throw e
+        }
+    }
+
+  }
+}

spark/src/test/scala/org/apache/comet/objectstore/NativeConfigSuite.scala

@@ -21,12 +21,23 @@ package org.apache.comet.objectstore
 
 import java.net.URI
 
+import scala.collection.mutable
+import scala.util.Try
+
+import org.scalatest.BeforeAndAfterEach
 import org.scalatest.funsuite.AnyFunSuite
 import org.scalatest.matchers.should.Matchers
 
 import org.apache.hadoop.conf.Configuration
 
-class NativeConfigSuite extends AnyFunSuite with Matchers {
+import org.apache.comet.CometNativeException
+import org.apache.comet.rules.CometScanRule
+
+class NativeConfigSuite extends AnyFunSuite with Matchers with BeforeAndAfterEach {
+
+  override protected def beforeEach(): Unit = {
+    CometScanRule.configValidityMap.clear()
+  }
 
   test("extractObjectStoreOptions - multiple cloud provider configurations") {
     val hadoopConf = new Configuration()
@@ -70,4 +81,71 @@ class NativeConfigSuite extends AnyFunSuite with Matchers {
       new URI("unsupported://test-bucket/test-object"))
     assert(unsupportedOptions.isEmpty, "Unsupported scheme should return empty options")
   }
+
+  test("validate object store config - no provider") {
+    val hadoopConf = new Configuration()
+    validate(hadoopConf)
+  }
+
+  test("validate object store config - valid providers") {
+    val hadoopConf = new Configuration()
+    val provider1 = "com.amazonaws.auth.EnvironmentVariableCredentialsProvider"
+    val provider2 = "com.amazonaws.auth.WebIdentityTokenCredentialsProvider"
+    hadoopConf.set("fs.s3a.aws.credentials.provider", Seq(provider1, provider2).mkString(","))
+    validate(hadoopConf)
+  }
+
+  test("validate object store config - invalid provider") {
+    val hadoopConf = new Configuration()
+    hadoopConf.set("fs.s3a.aws.credentials.provider", "invalid")
+    val e = intercept[CometNativeException] {
+      validate(hadoopConf)
+    }
+    assert(e.getMessage.contains("Unsupported credential provider: invalid"))
+  }
+
+  test("validate object store config - mixed anonymous providers") {
+    val hadoopConf = new Configuration()
+    val provider1 = "com.amazonaws.auth.AnonymousAWSCredentials"
+    val provider2 = "software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider"
+    hadoopConf.set("fs.s3a.aws.credentials.provider", Seq(provider1, provider2).mkString(","))
+    val e = intercept[CometNativeException] {
+      validate(hadoopConf)
+    }
+    val expectedError =
+      "Anonymous credential provider cannot be mixed with other credential providers"
+    assert(e.getMessage.contains(expectedError))
+  }
+
+  test("validate object store config - custom s3 endpoint not supported") {
+    val hadoopConf = new Configuration()
+    hadoopConf.set("fs.s3a.endpoint", "https://acme.storage.com")
+    val e = intercept[CometNativeException] {
+      validate(hadoopConf)
+    }
+    val expectedError =
+      "Custom S3 endpoints are not supported"
+    assert(e.getMessage.contains(expectedError))
+  }
+
+  test("validity cache") {
+    val hadoopConf = new Configuration()
+    hadoopConf.set("fs.s3a.endpoint", "https://acme.storage.com")
+
+    assert(CometScanRule.configValidityMap.isEmpty)
+    for (_ <- 0 until 5) {
+      assert(Try(validate(hadoopConf)).isFailure)
+      assert(CometScanRule.configValidityMap.size == 1)
+    }
+
+    hadoopConf.set("fs.s3a.endpoint", "https://acme2.storage.com")
+    assert(Try(validate(hadoopConf)).isFailure)
+    assert(CometScanRule.configValidityMap.size == 2)
+  }
+
+  private def validate(hadoopConf: Configuration): Unit = {
+    val path = "s3a://path/to/file.parquet"
+    val fallbackReasons = mutable.ListBuffer[String]()
+    CometScanRule.validateObjectStoreConfig(path, hadoopConf, fallbackReasons)
+  }
 }