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) + } + } +}