From b93bd7e6c95e1c9f9e59ea970e3c8f35e1fe96a2 Mon Sep 17 00:00:00 2001 From: bvolpato Date: Thu, 3 Sep 2026 19:22:39 +0000 Subject: [PATCH] fix(spark)!: preserve partition values in file scans A partitioned Parquet read containing rows (1,10) and (2,20) round-trips as (1,NULL) and (2,NULL). The exporter lists leaf files with the full table schema, while the importer treats virtual partition columns as physical file columns. Export each selected partition as a scan of its data schema with typed partition literals, preserve the merged column order through emit, and combine partitions with UNION_ALL. This uses Spark's resolved values rather than reconstructing them from paths, including nulls, explicit types, overlapping columns, pruned indexes, and multiple roots. Serialize canonical file URIs and decode them as URIs so escaped partition values address the original files. Local paths with unescaped spaces remain accepted. The representation uses standard Substrait v0.102.0 relations and adds one scan/project branch per nonempty partition. BREAKING CHANGE: Spark 3.4 conversion now rejects partitioned file scans whose file and partition column names differ only by case under case-insensitive resolution. Resolve the schema collision before conversion; this prevents changing the source plan's values to the different partition values. --- .../substrait/spark/compat/SparkCompat.scala | 3 + .../spark/logical/ToLogicalPlan.scala | 13 +- .../spark/logical/ToSubstraitRel.scala | 95 +++++++- .../spark/compat/SparkCompatImpl.scala | 2 + .../spark/PartitionedFilesSuite.scala | 226 ++++++++++++++++++ 5 files changed, 334 insertions(+), 5 deletions(-) create mode 100644 spark/src/test/scala/io/substrait/spark/PartitionedFilesSuite.scala diff --git a/spark/src/main/scala/io/substrait/spark/compat/SparkCompat.scala b/spark/src/main/scala/io/substrait/spark/compat/SparkCompat.scala index 418234671..49f028dc9 100644 --- a/spark/src/main/scala/io/substrait/spark/compat/SparkCompat.scala +++ b/spark/src/main/scala/io/substrait/spark/compat/SparkCompat.scala @@ -10,6 +10,9 @@ import org.apache.spark.sql.execution.datasources.{HadoopFsRelation, LogicalRela */ trait SparkCompat { + /** Whether partition values override differently cased file columns in case-insensitive mode. */ + def supportsCaseInsensitivePartitionOverlap: Boolean = true + /** Create a LogicalRelation with version-appropriate constructor */ def createLogicalRelation( relation: HadoopFsRelation, diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala index c3fa30dba..1c8734983 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToLogicalPlan.scala @@ -53,7 +53,7 @@ import io.substrait.relation.physical.{BroadcastExchange, MultiBucketExchange, R import io.substrait.util.EmptyVisitationContext import org.apache.hadoop.fs.Path -import java.net.URI +import java.net.{URI, URISyntaxException} import java.util.Optional import scala.annotation.nowarn @@ -440,7 +440,7 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes val (format, options) = convertFileFormat(formats.head) val location = SparkCompat.instance.createInMemoryFileIndex( spark, - localFiles.getItems.asScala.map(i => new Path(i.getPath.get())).toSeq, + localFiles.getItems.asScala.map(i => toFilePath(i.getPath.get())).toSeq, Map(), Some(schema)) val hadoopFsRelation = SparkCompat.instance.createHadoopFsRelation( @@ -461,6 +461,15 @@ class ToLogicalPlan(val spark: AnyRef = SparkCompat.instance.getOrCreateSparkSes remap(plan, localFiles.getRemap) } + private def toFilePath(path: String): Path = { + try { + new Path(new URI(path)) + } catch { + // Preserve support for unescaped local paths, such as filenames containing spaces. + case _: URISyntaxException => new Path(path) + } + } + def convertFileFormat(fileFormat: FileFormat): (SparkFileFormat, Map[String, String]) = { fileFormat match { case csv: FileFormat.DelimiterSeparatedTextReadOptions => diff --git a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala index 90f4a3538..6e9543c73 100644 --- a/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala +++ b/spark/src/main/scala/io/substrait/spark/logical/ToSubstraitRel.scala @@ -17,14 +17,14 @@ package io.substrait.spark.logical import io.substrait.spark.{FileHolder, SparkExtension, ToSubstraitType} -import io.substrait.spark.compat.WindowGroupLimitCase +import io.substrait.spark.compat.{SparkCompat, WindowGroupLimitCase} import io.substrait.spark.expression._ import io.substrait.spark.utils.Util import org.apache.spark.internal.Logging import org.apache.spark.sql.SaveMode import org.apache.spark.sql.catalyst.InternalRow -import org.apache.spark.sql.catalyst.analysis.ResolvedIdentifier +import org.apache.spark.sql.catalyst.analysis.{caseInsensitiveResolution, caseSensitiveResolution, ResolvedIdentifier} import org.apache.spark.sql.catalyst.catalog.{CatalogTable, HiveTableRelation} import org.apache.spark.sql.catalyst.expressions._ import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Average, Sum} @@ -38,6 +38,7 @@ import org.apache.spark.sql.execution.datasources.orc.OrcFileFormat import org.apache.spark.sql.execution.datasources.parquet.ParquetFileFormat import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation, DataSourceV2ScanRelation, V2SessionCatalog} import org.apache.spark.sql.hive.execution.{CreateHiveTableAsSelectCommand, InsertIntoHiveTable} +import org.apache.spark.sql.internal.SQLConf import org.apache.spark.sql.types.{NullType, StructField, StructType} import io.substrait.`type`.{NamedStruct, Type} @@ -511,6 +512,92 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { .build() } + private def buildPartitionedFileScan(fsRelation: HadoopFsRelation): relation.Rel = { + // LocalFiles does not carry directory partition values. Encode Spark's resolved values as + // literals so consumers do not have to infer their types or base paths again. + val fileFormat = convertFileFormat(fsRelation.fileFormat, fsRelation.options) + val dataSchema = ToSubstraitType.toNamedStruct(fsRelation.dataSchema) + val resolver = + if ( + SparkCompat.instance.getConf(fsRelation.sparkSession, SQLConf.CASE_SENSITIVE.key).toBoolean + ) { + caseSensitiveResolution + } else { + caseInsensitiveResolution + } + + if ( + !SparkCompat.instance.supportsCaseInsensitivePartitionOverlap && + fsRelation.dataSchema.exists( + data => + fsRelation.partitionSchema.exists( + partition => data.name != partition.name && resolver(data.name, partition.name))) + ) { + throw new UnsupportedOperationException( + "This Spark version cannot reliably read file and partition columns that differ only in case") + } + + // Partition values override overlapping file columns, preserving the merged schema's order. + val outputMapping = fsRelation.schema.map { + field => + val partitionIndex = + fsRelation.partitionSchema.indexWhere(p => resolver(p.name, field.name)) + if (partitionIndex >= 0) { + fsRelation.dataSchema.size + partitionIndex + } else { + fsRelation.dataSchema.indexWhere(d => resolver(d.name, field.name)) + } + } + val remap = relation.Rel.Remap.of(outputMapping.map(Int.box).toSeq.asJava) + + val partitions = fsRelation.location.listFiles(Nil, Nil).filter(_.files.nonEmpty).map { + partition => + val read = relation.LocalFiles + .builder() + .initialSchema(dataSchema) + .addAllItems( + partition.files + .map { + file => + FileOrFiles + .builder() + .fileFormat(fileFormat) + .partitionIndex(0) + .start(0) + .length(file.getLen) + .path(file.getPath.toUri.toString) + .pathType(PathType.URI_FILE) + .build() + } + .toSeq + .asJava) + .build() + val values = fsRelation.partitionSchema.zipWithIndex.map { + case (field, index) => + ToSubstraitLiteral( + Literal(partition.values.get(index, field.dataType), field.dataType), + Some(field.nullable)) + } + relation.Project + .builder() + .input(read) + .addAllExpressions(values.toSeq.asJava) + .remap(remap) + .build() + } + + partitions.size match { + case 0 => + relation.VirtualTableScan + .builder() + .initialSchema(ToSubstraitType.toNamedStruct(fsRelation.schema)) + .build() + case 1 => partitions.head + case _ => + relation.Set.builder().setOp(SetOp.UNION_ALL).addAllInputs(partitions.toSeq.asJava).build() + } + } + private def convertFileFormat( fileFormat: DSFileFormat, options: Map[String, String]): FileFormat = fileFormat match { @@ -533,7 +620,7 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { } /** Read Operator: https://substrait.io/relations/logical_relations/#read-operator */ - private def convertReadOperator(plan: LeafNode): relation.AbstractReadRel = { + private def convertReadOperator(plan: LeafNode): relation.Rel = { var tableNames: List[String] = null plan match { case logicalRelation: LogicalRelation if logicalRelation.catalogTable.isDefined => @@ -558,6 +645,8 @@ class ToSubstraitRel extends AbstractLogicalPlanVisitor with Logging { buildVirtualTableScan(rdd.schema, rdd.rdd.take(_rddLimit).toIndexedSeq) case logicalRelation: LogicalRelation => logicalRelation.relation match { + case fsRelation: HadoopFsRelation if fsRelation.partitionSchema.nonEmpty => + buildPartitionedFileScan(fsRelation) case fsRelation: HadoopFsRelation => buildLocalFileScan(fsRelation) case _ => diff --git a/spark/src/main/spark-3.4/io/substrait/spark/compat/SparkCompatImpl.scala b/spark/src/main/spark-3.4/io/substrait/spark/compat/SparkCompatImpl.scala index 022882eee..9c41738f1 100644 --- a/spark/src/main/spark-3.4/io/substrait/spark/compat/SparkCompatImpl.scala +++ b/spark/src/main/spark-3.4/io/substrait/spark/compat/SparkCompatImpl.scala @@ -8,6 +8,8 @@ import io.substrait.relation class SparkCompatImpl extends SparkCompat { + override def supportsCaseInsensitivePartitionOverlap: Boolean = false + override def createLogicalRelation( relation: HadoopFsRelation, output: Seq[AttributeReference], diff --git a/spark/src/test/scala/io/substrait/spark/PartitionedFilesSuite.scala b/spark/src/test/scala/io/substrait/spark/PartitionedFilesSuite.scala new file mode 100644 index 000000000..8a916fff3 --- /dev/null +++ b/spark/src/test/scala/io/substrait/spark/PartitionedFilesSuite.scala @@ -0,0 +1,226 @@ +package io.substrait.spark + +import io.substrait.spark.compat.SparkCompat +import io.substrait.spark.logical.{ToLogicalPlan, ToSubstraitRel} + +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.analysis.caseSensitiveResolution +import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan +import org.apache.spark.sql.classic.DatasetUtil +import org.apache.spark.sql.execution.datasources.{FileIndex, HadoopFsRelation, LogicalRelation, PartitionDirectory} +import org.apache.spark.sql.test.SharedSparkSession +import org.apache.spark.sql.types.{DataType, IntegerType, LongType, StringType, StructField, StructType} + +import io.substrait.plan.{PlanProtoConverter, ProtoPlanConverter} +import io.substrait.relation.{LocalFiles => SubstraitLocalFiles} +import io.substrait.relation.files.FileOrFiles +import org.apache.hadoop.fs.Path + +import java.net.URI +import java.time.LocalDate + +import scala.jdk.CollectionConverters._ + +class PartitionedFilesSuite extends SharedSparkSession { + + private def assertRoundTrip(plan: LogicalPlan, expected: Seq[Row]): Unit = { + val original = DatasetUtil.fromLogicalPlan(spark, plan).collect().toSeq + assertResult(expected.sortBy(_.toString))(original.sortBy(_.toString)) + + val substrait = new ToSubstraitRel().convert(plan) + val bytes = new PlanProtoConverter().toProto(substrait).toByteArray + val decoded = new ProtoPlanConverter().from(io.substrait.proto.Plan.parseFrom(bytes)) + assertResult(substrait)(decoded) + + val converted = new ToLogicalPlan(spark).convert(decoded) + assert( + DataType.equalsStructurallyByName(plan.schema, converted.schema, caseSensitiveResolution)) + val actual = DatasetUtil.fromLogicalPlan(spark, converted).collect().toSeq + assertResult(expected.sortBy(_.toString))(actual.sortBy(_.toString)) + } + + Seq("parquet", "orc", "csv").foreach { + format => + test(s"partition values survive $format reads and file options") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark + .sql("select 1 id, 'left|right' value, 10 part union all select 2, 'other', 20") + .write + .format(format) + .option("header", true) + .option("delimiter", "|") + .partitionBy("part") + .save(path) + val schema = StructType( + Seq( + StructField("id", IntegerType), + StructField("value", StringType), + StructField("part", IntegerType))) + val data = spark.read + .format(format) + .schema(schema) + .option("header", true) + .option("delimiter", "|") + .load(path) + assertRoundTrip( + data.queryExecution.optimizedPlan, + Seq(Row(1, "left|right", 10), Row(2, "other", 20))) + } + } + } + + test("multiple roots and basePath preserve the selected partition values") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark + .sql("select 1 id, 10 part union all select 2, 20 union all select 3, 30") + .write + .partitionBy("part") + .parquet(path) + val selected = spark.read + .option("basePath", path) + .parquet(s"$path/part=10", s"$path/part=20") + assertRoundTrip(selected.queryExecution.optimizedPlan, Seq(Row(1, 10), Row(2, 20))) + assertRoundTrip(selected.filter("part = 10").queryExecution.optimizedPlan, Seq(Row(1, 10))) + } + } + + test("date null and escaped string partition values retain their types and values") { + withSQLConf("spark.sql.datetime.java8API.enabled" -> "true") { + withTempPath { + directory => + val path = directory.getAbsolutePath + "/root with spaces" + spark + .sql("select 1 id, date '2024-01-02' day, 'a/b% c' label " + + "union all select 2, cast(null as date), cast(null as string)") + .write + .partitionBy("day", "label") + .parquet(path) + val data = spark.read.parquet(path) + assertRoundTrip( + data.queryExecution.optimizedPlan, + Seq(Row(1, LocalDate.of(2024, 1, 2), "a/b% c"), Row(2, null, null))) + } + } + } + + test("local file reads still accept unescaped paths containing spaces") { + withTempPath { + directory => + val path = directory.getAbsolutePath + "/root with spaces" + spark.sql("select 1 id").write.parquet(path) + val original = spark.read.parquet(path).queryExecution.optimizedPlan + val scan = new ToSubstraitRel().visit(original).asInstanceOf[SubstraitLocalFiles] + val files = scan.getItems.asScala.map { + file => FileOrFiles.builder().from(file).path(new URI(file.getPath.get()).getPath).build() + } + val rawPaths = SubstraitLocalFiles.builder().from(scan).items(files.toSeq.asJava).build() + val converted = new ToLogicalPlan(spark).convert(rawPaths) + assertResult(Seq(Row(1)))(DatasetUtil.fromLogicalPlan(spark, converted).collect().toSeq) + } + } + + test("explicit partition types are retained") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark.sql("select 1 id, 10 part").write.partitionBy("part").parquet(path) + val schema = StructType(Seq(StructField("id", IntegerType), StructField("part", LongType))) + val data = spark.read.schema(schema).option("basePath", path).parquet(s"$path/part=10") + assertRoundTrip(data.queryExecution.optimizedPlan, Seq(Row(1, 10L))) + } + } + + test("partition values override overlapping file columns in merged schema order") { + withSQLConf("spark.sql.caseSensitive" -> "false") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark.sql("select 1 id, '999' p, 'physical' value").write.parquet(s"$path/p=10") + val data = spark.read.parquet(path) + assertResult(Seq("id", "p", "value"))(data.columns.toSeq) + assertRoundTrip(data.queryExecution.optimizedPlan, Seq(Row(1, 10, "physical"))) + } + } + } + + test("mixed-case overlapping columns are rejected when Spark does not override them") { + withSQLConf("spark.sql.caseSensitive" -> "false") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark.sql("select 1 id, 999 P, 'physical' value").write.parquet(s"$path/p=10") + val data = spark.read.parquet(path) + val plan = data.queryExecution.optimizedPlan + if (SparkCompat.instance.supportsCaseInsensitivePartitionOverlap) { + assertRoundTrip(plan, Seq(Row(1, 10, "physical"))) + } else { + assertResult(Seq(Row(1, 999, "physical")))(data.collect().toSeq) + val error = intercept[UnsupportedOperationException] { + new ToSubstraitRel().convert(plan) + } + assert(error.getMessage.contains("differ only in case")) + } + } + } + } + + test("case-sensitive file and partition column names remain distinct") { + withSQLConf("spark.sql.caseSensitive" -> "true") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark.sql("select 1 id, 999 P, 'physical' value").write.parquet(s"$path/p=10") + val data = spark.read.parquet(path) + assertRoundTrip(data.queryExecution.optimizedPlan, Seq(Row(1, 999, "physical", 10))) + } + } + } + + private def withPartitions( + original: HadoopFsRelation, + partitions: Seq[PartitionDirectory]): LogicalPlan = { + val index = new FileIndex { + override def rootPaths: Seq[Path] = original.location.rootPaths + override def listFiles( + partitionFilters: Seq[Expression], + dataFilters: Seq[Expression]): Seq[PartitionDirectory] = partitions + override def inputFiles: Array[String] = + partitions.flatMap(_.files.map(_.getPath.toString)).toArray + override def refresh(): Unit = () + override def sizeInBytes: Long = partitions.flatMap(_.files.map(_.getLen)).sum + override def partitionSchema: StructType = original.partitionSchema + } + val relation = original.copy(location = index)(spark) + SparkCompat.instance.createLogicalRelation( + relation, + ToSparkType.toAttributeSeq(ToSubstraitType.toNamedStruct(relation.schema)), + None, + false) + } + + test("pruned and empty file indexes do not restore excluded partitions") { + withTempPath { + directory => + val path = directory.getAbsolutePath + spark + .sql("select 1 id, 10 part union all select 2, 20") + .write + .partitionBy("part") + .parquet(path) + val logical = spark.read + .parquet(path) + .queryExecution + .optimizedPlan + .asInstanceOf[LogicalRelation] + val original = logical.relation.asInstanceOf[HadoopFsRelation] + val selected = original.location.listFiles(Nil, Nil).filter(_.values.getInt(0) == 10) + assertRoundTrip(withPartitions(original, selected), Seq(Row(1, 10))) + assertRoundTrip(withPartitions(original, Seq.empty), Seq.empty) + } + } +}