[SPARK-58609][SQL] Add archivePathFilter option to select inner archive entries

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/FileSourceOptions.scala

@@ -18,7 +18,11 @@ package org.apache.spark.sql.catalyst
 
 import java.util.regex.{Pattern, PatternSyntaxException}
 
-import org.apache.spark.sql.catalyst.FileSourceOptions.{IGNORE_CORRUPT_FILES, IGNORE_MISSING_FILES, IGNORED_PATH_SEGMENT_REGEX}
+import scala.util.control.NonFatal
+
+import org.apache.hadoop.fs.GlobPattern
+
+import org.apache.spark.sql.catalyst.FileSourceOptions.{ARCHIVE_PATH_FILTER, IGNORE_CORRUPT_FILES, IGNORE_MISSING_FILES, IGNORED_PATH_SEGMENT_REGEX}
 import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateFormatter}
 import org.apache.spark.sql.internal.{LegacyBehaviorPolicy, SQLConf}
 
@@ -66,12 +70,37 @@ class FileSourceOptions(
    */
   lazy val ignoredPathSegmentRegexPattern: Pattern =
     FileSourceOptions.compileIgnoredPathSegmentRegex(ignoredPathSegmentRegex)
+
+  /**
+   * Glob selecting which inner archive entries to read, matched against each entry's full path
+   * within the archive (e.g. `subdir/*`, `*/*.csv`). An empty value disables the filter, matching
+   * how an empty [[ignoredPathSegmentRegex]] is treated. Validated here so an invalid glob fails on
+   * the driver.
+   */
+  val archivePathFilter: Option[String] = {
+    val glob = parameters.get(ARCHIVE_PATH_FILTER).filter(_.nonEmpty)
+    glob.foreach(FileSourceOptions.compileArchivePathFilter)
+    glob
+  }
+
+  /**
+   * The effective [[archivePathFilter]] is compiled once per instance of this class, so the archive
+   * reads sharing one options object reuse a single matcher rather than re-compiling per archive.
+   * `transient` because Hadoop's `GlobPattern` is not serializable, so an executor recompiles from
+   * [[archivePathFilter]] on first use. Paths that cannot reach this value -- the schema-inference
+   * RDDs, which would have to capture the matcher in a closure, and the parallel footer/schema
+   * readers, whose signatures are fixed -- carry [[archivePathFilter]] instead and compile it
+   * themselves.
+   */
+  @transient lazy val archivePathFilterPattern: Option[GlobPattern] =
+    archivePathFilter.map(FileSourceOptions.compileArchivePathFilter)
 }
 
 object FileSourceOptions {
   val IGNORE_CORRUPT_FILES = "ignoreCorruptFiles"
   val IGNORE_MISSING_FILES = "ignoreMissingFiles"
   val IGNORED_PATH_SEGMENT_REGEX = "ignoredPathSegmentRegex"
+  val ARCHIVE_PATH_FILTER = "archivePathFilter"
 
   // A regex that never matches any name, used when the filter is disabled by an empty value.
   private val DISABLED_FILTER_PATTERN = Pattern.compile("(?!)")
@@ -97,4 +126,15 @@ object FileSourceOptions {
       }
     }
   }
+
+  /** Compiles `archivePathFilter` into a glob matcher, reporting an invalid glob clearly. */
+  def compileArchivePathFilter(glob: String): GlobPattern = {
+    try {
+      new GlobPattern(glob)
+    } catch {
+      case NonFatal(e) =>
+        throw new IllegalArgumentException(
+          s"The '$ARCHIVE_PATH_FILTER' value '$glob' is not a valid glob.", e)
+    }
+  }
 }

sql/core/src/main/scala/org/apache/spark/sql/avro/AvroFileFormat.scala

@@ -199,7 +199,9 @@ private[sql] class AvroFileFormat extends FileFormat
     } else {
       new NoopFilters
     }
-    SupportsArchiveFormat.readArchiveEntries(file.toPath, conf) { (_, in) =>
+    val entryGlob = parsedOptions.archivePathFilterPattern
+    SupportsArchiveFormat.readArchiveEntries(
+        file.toPath, conf, archivePathFilter = entryGlob) { (_, in) =>
       val datumReader = userProvidedSchema match {
         case Some(schema) => new GenericDatumReader[GenericRecord](schema)
         case None => new GenericDatumReader[GenericRecord]()

sql/core/src/main/scala/org/apache/spark/sql/avro/AvroUtils.scala

@@ -27,7 +27,7 @@ import org.apache.avro.generic.{GenericDatumReader, GenericRecord}
 import org.apache.avro.mapred.{AvroOutputFormat, FsInput}
 import org.apache.avro.mapreduce.AvroJob
 import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{FileStatus, Path}
+import org.apache.hadoop.fs.{FileStatus, GlobPattern, Path}
 import org.apache.hadoop.mapreduce.Job
 
 import org.apache.spark.{SparkException, SparkIllegalArgumentException}
@@ -97,7 +97,8 @@ private[sql] object AvroUtils extends Logging {
         }
         if (archives.nonEmpty) {
           inferAvroSchemaFromArchives(archives, nonArchives, conf, parsedOptions.ignoreExtension,
-            fileSourceOptions.ignoreCorruptFiles, fileSourceOptions.ignoreMissingFiles)
+            fileSourceOptions.ignoreCorruptFiles, fileSourceOptions.ignoreMissingFiles,
+            fileSourceOptions.archivePathFilterPattern)
         } else {
           inferAvroSchemaFromFiles(files, conf, parsedOptions.ignoreExtension,
             fileSourceOptions.ignoreCorruptFiles)
@@ -250,10 +251,12 @@ private[sql] object AvroUtils extends Logging {
       conf: Configuration,
       ignoreExtension: Boolean,
       ignoreCorruptFiles: Boolean,
-      ignoreMissingFiles: Boolean): Schema = {
+      ignoreMissingFiles: Boolean,
+      archivePathFilter: Option[GlobPattern]): Schema = {
     archives.iterator
       .flatMap { f =>
-        firstArchiveEntrySchema(f.getPath, conf, ignoreCorruptFiles, ignoreMissingFiles)
+        firstArchiveEntrySchema(
+          f.getPath, conf, ignoreCorruptFiles, ignoreMissingFiles, archivePathFilter)
       }
       .nextOption()
       .getOrElse {
@@ -274,11 +277,13 @@ private[sql] object AvroUtils extends Logging {
       path: Path,
       conf: Configuration,
       ignoreCorruptFiles: Boolean,
-      ignoreMissingFiles: Boolean): Option[Schema] = {
+      ignoreMissingFiles: Boolean,
+      archivePathFilter: Option[GlobPattern]): Option[Schema] = {
     try {
       // `readArchiveEntries` returns a Closeable iterator; take the first entry's schema and close
       // it so the archive stream is released without draining the remaining entries.
-      val entries = SupportsArchiveFormat.readArchiveEntries(path, conf) { (_, in) =>
+      val entries = SupportsArchiveFormat.readArchiveEntries(
+          path, conf, archivePathFilter = archivePathFilter) { (_, in) =>
         val stream = new DataFileStream[GenericRecord](in, new GenericDatumReader[GenericRecord]())
         try {
           Iterator.single(stream.getSchema)

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/SupportsArchiveFormat.scala

@@ -34,7 +34,7 @@ import org.apache.commons.compress.archivers.zip.ZipFile
 import org.apache.commons.io.ByteOrderMark
 import org.apache.commons.io.input.{BOMInputStream, CloseShieldInputStream}
 import org.apache.hadoop.conf.Configuration
-import org.apache.hadoop.fs.{FSDataInputStream, Path}
+import org.apache.hadoop.fs.{FSDataInputStream, GlobPattern, Path}
 import org.apache.hadoop.io.Text
 import org.apache.hadoop.util.LineReader
 
@@ -104,16 +104,18 @@ trait SupportsArchiveFormat extends Logging {
    * Reads an archive by unpacking each entry to a temp file and applying `readEntry`, for a
    * format that needs a complete file on disk (random access).
    *
-   * @param file       the archive as a [[PartitionedFile]]
-   * @param conf       Hadoop configuration used to open the archive
-   * @param tempPrefix prefix for the per-task temp dir the entries are unpacked into
-   * @param readEntry  reads one unpacked entry file into rows
+   * @param file              the archive as a [[PartitionedFile]]
+   * @param conf              Hadoop configuration used to open the archive
+   * @param tempPrefix        prefix for the per-task temp dir the entries are unpacked into
+   * @param archivePathFilter optional glob matched against the entry's full path
+   * @param readEntry         reads one unpacked entry file into rows
    * @return iterator of rows across all entries
    */
   protected def readLocalizedEntries(
       file: PartitionedFile,
       conf: Configuration,
-      tempPrefix: String)(
+      tempPrefix: String,
+      archivePathFilter: Option[GlobPattern])(
       readEntry: PartitionedFile => Iterator[InternalRow]): Iterator[InternalRow] = {
     val tempDir = Utils.createTempDir(Utils.getLocalDir(SparkEnv.get.conf), tempPrefix)
     // Register cleanup before constructing `entries`, which can throw before returning an iterator
@@ -122,7 +124,8 @@ trait SupportsArchiveFormat extends Logging {
       Utils.deleteRecursively(tempDir)
     })
     val entries =
-      try SupportsArchiveFormat.localizeEntries(file.toPath, conf, tempDir, archiveEntryFilter)
+      try SupportsArchiveFormat.localizeEntries(
+        file.toPath, conf, tempDir, archiveEntryFilter, archivePathFilter)
       catch {
         case NonFatal(e) =>
           Utils.deleteRecursively(tempDir)
@@ -342,10 +345,16 @@ object SupportsArchiveFormat {
    *
    * @param entry                   the archive entry to test
    * @param ignoredPathSegmentRegex per-segment filter matched against each `/`-separated component
-   * @return true if the entry is a directory or any path component is filtered out
+   * @param archivePathFilter       optional glob matched against the entry's full path
+   * @return true if the entry is a directory, any path component is filtered out, or the entry's
+   *         path does not match `archivePathFilter`
    */
-  private def shouldSkipEntry(entry: ArchiveEntry, ignoredPathSegmentRegex: Pattern): Boolean = {
+  private def shouldSkipEntry(
+      entry: ArchiveEntry,
+      ignoredPathSegmentRegex: Pattern,
+      archivePathFilter: Option[GlobPattern]): Boolean = {
     if (entry.isDirectory) return true
+    if (archivePathFilter.exists(!_.matches(entry.getName))) return true
     entry.getName.split("/").exists(c =>
       c.nonEmpty && HadoopFSUtils.shouldFilterOutPathName(c, ignoredPathSegmentRegex))
   }
@@ -359,13 +368,15 @@ object SupportsArchiveFormat {
    * @param ignoredPathSegmentRegex per-segment filter for entries to skip (defaults to the
    *                                `InMemoryFileIndex` filter); pass a custom one to match a
    *                                loose-file scan
+   * @param archivePathFilter       optional glob matched against the entry's full path
    * @param parseEntry              turns one entry's `(entry, stream)` into an iterator of results
    * @return the concatenated results across kept entries, lazily one entry at a time
    */
   def readArchiveEntries[T](
       path: Path,
       conf: Configuration,
-      ignoredPathSegmentRegex: Pattern = HadoopFSUtils.defaultIgnoredPathSegmentRegexPattern)(
+      ignoredPathSegmentRegex: Pattern = HadoopFSUtils.defaultIgnoredPathSegmentRegexPattern,
+      archivePathFilter: Option[GlobPattern])(
       parseEntry: (ArchiveEntry, InputStream) => Iterator[T]): Iterator[T] = {
     val archive = openArchiveStream(path, conf)
     var closed = false
@@ -395,7 +406,9 @@ object SupportsArchiveFormat {
           var next: (ArchiveEntry, InputStream) = null
           while (next == null && archive.hasNext) {
             val entry = archive.next()
-            if (!shouldSkipEntry(entry._1, ignoredPathSegmentRegex)) next = entry
+            if (!shouldSkipEntry(entry._1, ignoredPathSegmentRegex, archivePathFilter)) {
+              next = entry
+            }
           }
           if (next == null) {
             done = true
@@ -460,17 +473,19 @@ object SupportsArchiveFormat {
    * companion so executor-side callers (a format's distributed archive inference) can use it
    * without a trait instance.
    *
-   * @param path        the archive path
-   * @param conf        Hadoop configuration used to open the archive
-   * @param localDir    directory the per-entry temp files are created under
-   * @param entryFilter which entry names to keep
+   * @param path              the archive path
+   * @param conf              Hadoop configuration used to open the archive
+   * @param localDir          directory the per-entry temp files are created under
+   * @param entryFilter       which entry names to keep
+   * @param archivePathFilter optional glob matched against the entry's full path
    */
   def localizeEntries(
       path: Path,
       conf: Configuration,
       localDir: File,
-      entryFilter: String => Boolean): Iterator[(String, File)] =
-    readArchiveEntries(path, conf) { (entry, in) =>
+      entryFilter: String => Boolean,
+      archivePathFilter: Option[GlobPattern]): Iterator[(String, File)] =
+    readArchiveEntries(path, conf, archivePathFilter = archivePathFilter) { (entry, in) =>
       val name = entry.getName
       if (entryFilter(name)) {
         Iterator.single((name, copyEntryToLocalFile(in, localDir, name)))

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/binaryfile/BinaryFileFormat.scala

@@ -113,16 +113,17 @@ case class BinaryFileFormat() extends FileFormat
     val caseInsensitiveOptions = CaseInsensitiveMap(options)
     val archiveReadEnabled = !caseInsensitiveOptions.get(WHOLE_FILE).forall(_.toBoolean) &&
       getSqlConf(sparkSession).getConf(SQLConf.ARCHIVE_FORMAT_READER_ENABLED)
+    val fileSourceOptions = new FileSourceOptions(caseInsensitiveOptions)
 
     file: PartitionedFile => {
       val path = file.toPath
       val fs = path.getFileSystem(broadcastedHadoopConf.value.value)
       val status = fs.getFileStatus(path)
       if (archiveReadEnabled && SupportsArchiveFormat.isArchivePath(path)) {
-        val ignoredPathSegmentRegex =
-          new FileSourceOptions(caseInsensitiveOptions).ignoredPathSegmentRegexPattern
         SupportsArchiveFormat.readArchiveEntries(
-            path, broadcastedHadoopConf.value.value, ignoredPathSegmentRegex) { (entry, in) =>
+            path, broadcastedHadoopConf.value.value,
+            fileSourceOptions.ignoredPathSegmentRegexPattern,
+            fileSourceOptions.archivePathFilterPattern) { (entry, in) =>
           val entryStatus = new FileStatus(
             entry.getSize, false, 0, 0, status.getModificationTime,
             new Path(s"${status.getPath}!/${entry.getName}"))