szehon-ho commented on code in PR #54459:
URL: https://github.com/apache/spark/pull/54459#discussion_r2893239467
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/v2/PushDownUtils.scala:
##########
@@ -22,24 +22,66 @@ import scala.collection.mutable
import org.apache.spark.sql.catalyst.expressions.{AttributeReference,
AttributeSet, Expression, NamedExpression, SchemaPruning}
import org.apache.spark.sql.catalyst.types.DataTypeUtils.toAttributes
import org.apache.spark.sql.catalyst.util.CharVarcharUtils
+import org.apache.spark.sql.connector.expressions.IdentityTransform
import org.apache.spark.sql.connector.expressions.SortOrder
-import org.apache.spark.sql.connector.expressions.filter.Predicate
+import org.apache.spark.sql.connector.expressions.filter.{PartitionPredicate,
Predicate}
import org.apache.spark.sql.connector.read.{Scan, ScanBuilder,
SupportsPushDownFilters, SupportsPushDownLimit, SupportsPushDownOffset,
SupportsPushDownRequiredColumns, SupportsPushDownTableSample,
SupportsPushDownTopN, SupportsPushDownV2Filters}
-import org.apache.spark.sql.execution.datasources.DataSourceStrategy
+import org.apache.spark.sql.execution.datasources.{DataSourceStrategy,
DataSourceUtils}
import org.apache.spark.sql.internal.SQLConf
-import org.apache.spark.sql.internal.connector.SupportsPushDownCatalystFilters
+import org.apache.spark.sql.internal.connector.{PartitionPredicateImpl,
SupportsPushDownCatalystFilters}
import org.apache.spark.sql.sources
-import org.apache.spark.sql.types.StructType
+import org.apache.spark.sql.types.{StructField, StructType}
import org.apache.spark.util.ArrayImplicits._
import org.apache.spark.util.collection.Utils
object PushDownUtils {
+
+ /**
+ * Returns partition schema as a StructType when the table partitioning.
+ * Currently only supported for identity transforms on simple (single-name)
field references.
+ *
+ * @return Some(StructType) for partition transform types, if supported.
+ */
+ def getPartitionSchemaForPartitionPredicate(
+ relation: DataSourceV2Relation): Option[StructType] = {
+ val partitioning = relation.table.partitioning().toIndexedSeq
+ val partitionColNamesOpt: Seq[Option[String]] = partitioning.map {
+ case id: IdentityTransform =>
+ id.ref.fieldNames().toIndexedSeq match {
+ case Seq(name) => Some(name)
+ case _ => None // Not supported for multiple field names (e.g.
nested field)
+ }
+ case _ => None
+ }
+ partitionColNamesOpt match {
+ // Only support identity transform on simple field reference
+ case seq if seq.isEmpty || seq.exists(_.isEmpty) => None
Review Comment:
This may not be do-able now. The contract is the DSV2 connector needs to
pass an `InternalRow partitionKeys` to PartitionPredicate.accept() that matches
the Table.partitioning() schema. Now we are saying, for the time being let
them pass one that only matches a pruned partition schema that only has
identity columns (so it doesnt match the actual Table.partitioning() schema) .
Once we fully support the PartitionPredicate for arbitrary partition
transform, we can revisit?
By the way, the only consumer I know yet is Iceberg, which use the partition
transform year(), hour(), etc
--
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]
