This is an automated email from the ASF dual-hosted git repository.
szehon-ho pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg.git
The following commit(s) were added to refs/heads/main by this push:
new 23774790e6 Spark 4.2: Push down opaque partition predicates (#17956)
23774790e6 is described below
commit 23774790e6928610e4e9ef7192373eb68b20b0e2
Author: Szehon Ho <[email protected]>
AuthorDate: Tue Sep 22 16:10:14 2026 -0700
Spark 4.2: Push down opaque partition predicates (#17956)
---
.../iceberg/spark/source/BaseSparkScanBuilder.java | 27 ++-
.../iceberg/spark/source/SparkBatchQueryScan.java | 15 +-
.../iceberg/spark/source/SparkCopyOnWriteScan.java | 15 +-
.../spark/source/SparkIncrementalAppendScan.java | 15 +-
.../spark/source/SparkPartitioningAwareScan.java | 125 +++++++++++++-
.../spark/source/SparkRuntimeFilterableScan.java | 13 +-
.../iceberg/spark/source/SparkScanBuilder.java | 8 +
.../iceberg/spark/source/StructInternalRow.java | 6 +-
.../iceberg/spark/source/TestFilteredScan.java | 90 +++++++++-
.../iceberg/spark/sql/TestFilterPushDown.java | 191 +++++++++++++++++++++
10 files changed, 490 insertions(+), 15 deletions(-)
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/BaseSparkScanBuilder.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/BaseSparkScanBuilder.java
index be34eaeed5..853c6059ae 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/BaseSparkScanBuilder.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/BaseSparkScanBuilder.java
@@ -44,6 +44,7 @@ import org.apache.iceberg.types.TypeUtil;
import org.apache.iceberg.types.TypeUtil.GetID;
import org.apache.iceberg.types.Types;
import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
import org.apache.spark.sql.connector.expressions.filter.Predicate;
import org.apache.spark.sql.connector.read.ScanBuilder;
import org.apache.spark.sql.types.StructField;
@@ -62,7 +63,6 @@ import org.slf4j.LoggerFactory;
abstract class BaseSparkScanBuilder implements ScanBuilder {
private static final Logger LOG =
LoggerFactory.getLogger(BaseSparkScanBuilder.class);
- private static final Predicate[] NO_PREDICATES = new Predicate[0];
private final SparkSession spark;
private final Table table;
@@ -71,10 +71,11 @@ abstract class BaseSparkScanBuilder implements ScanBuilder {
private final boolean caseSensitive;
private final Set<String> metaFieldNames = Sets.newLinkedHashSet();
private final InMemoryMetricsReporter metricsReporter = new
InMemoryMetricsReporter();
+ private final List<PartitionPredicate> partitionPredicates =
Lists.newArrayList();
+ private final List<Predicate> pushedPredicates = Lists.newArrayList();
+ private final List<Expression> filters = Lists.newArrayList();
private Schema projection;
- private List<Expression> filters = Lists.newArrayList();
- private Predicate[] pushedPredicates = NO_PREDICATES;
private Integer limit = null;
protected BaseSparkScanBuilder(
@@ -119,6 +120,10 @@ abstract class BaseSparkScanBuilder implements ScanBuilder
{
return filters.stream().reduce(Expressions.alwaysTrue(), Expressions::and);
}
+ protected List<PartitionPredicate> partitionPredicates() {
+ return partitionPredicates;
+ }
+
protected InMemoryMetricsReporter metricsReporter() {
return metricsReporter;
}
@@ -156,6 +161,16 @@ abstract class BaseSparkScanBuilder implements ScanBuilder
{
List<Predicate> postScanPredicates =
Lists.newArrayListWithExpectedSize(predicates.length);
for (Predicate predicate : predicates) {
+ if (predicate instanceof PartitionPredicate partitionPredicate) {
+ partitionPredicates.add(partitionPredicate);
+ pushablePredicates.add(predicate);
+
+ // TODO: Currently, all partition predicates are returned to Spark for
re-evaluation.
+ // Optimize by returning only predicates that cannot be evaluated
against all scanned specs.
+ postScanPredicates.add(predicate);
+ continue;
+ }
+
try {
Expression expr = SparkV2Filters.convert(predicate);
@@ -178,15 +193,15 @@ abstract class BaseSparkScanBuilder implements
ScanBuilder {
}
}
- this.filters = expressions;
- this.pushedPredicates = pushablePredicates.toArray(new Predicate[0]);
+ this.filters.addAll(expressions);
+ this.pushedPredicates.addAll(pushablePredicates);
return postScanPredicates.toArray(new Predicate[0]);
}
// logic necessary for SupportsPushDownV2Filters
public Predicate[] pushedPredicates() {
- return pushedPredicates;
+ return pushedPredicates.toArray(new Predicate[0]);
}
// logic necessary for SupportsPushDownLimit
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchQueryScan.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchQueryScan.java
index 814ca410d1..a7bfa9ecc0 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchQueryScan.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkBatchQueryScan.java
@@ -31,6 +31,7 @@ import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.metrics.ScanReport;
import org.apache.iceberg.spark.SparkReadConf;
import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
import org.apache.spark.sql.connector.read.Statistics;
class SparkBatchQueryScan extends SparkRuntimeFilterableScan {
@@ -48,8 +49,18 @@ class SparkBatchQueryScan extends SparkRuntimeFilterableScan
{
SparkReadConf readConf,
Schema projection,
List<Expression> filters,
+ List<PartitionPredicate> partitionPredicates,
Supplier<ScanReport> scanReportSupplier) {
- super(spark, table, schema, scan, readConf, projection, filters,
scanReportSupplier);
+ super(
+ spark,
+ table,
+ schema,
+ scan,
+ readConf,
+ projection,
+ filters,
+ partitionPredicates,
+ scanReportSupplier);
this.snapshot = snapshot;
this.branch = branch;
}
@@ -80,6 +91,7 @@ class SparkBatchQueryScan extends SparkRuntimeFilterableScan {
&& Objects.equals(branch, that.branch)
&& readSchema().equals(that.readSchema()) // compare Spark schemas to
ignore field ids
&& filtersDesc().equals(that.filtersDesc())
+ && partitionPredicates().equals(that.partitionPredicates())
&& runtimeFiltersDesc().equals(that.runtimeFiltersDesc());
}
@@ -92,6 +104,7 @@ class SparkBatchQueryScan extends SparkRuntimeFilterableScan
{
branch,
readSchema(),
filtersDesc(),
+ partitionPredicates(),
runtimeFiltersDesc());
}
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
index b8b3a0edea..528409f96b 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkCopyOnWriteScan.java
@@ -37,6 +37,7 @@ import org.apache.iceberg.spark.SparkReadConf;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.connector.expressions.Literal;
import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
import org.apache.spark.sql.connector.expressions.filter.Predicate;
import org.apache.spark.sql.connector.read.Statistics;
import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering;
@@ -63,8 +64,18 @@ class SparkCopyOnWriteScan extends
SparkPartitioningAwareScan<FileScanTask>
SparkReadConf readConf,
Schema projection,
List<Expression> filters,
+ List<PartitionPredicate> partitionPredicates,
Supplier<ScanReport> scanReportSupplier) {
- super(spark, table, schema, scan, readConf, projection, filters,
scanReportSupplier);
+ super(
+ spark,
+ table,
+ schema,
+ scan,
+ readConf,
+ projection,
+ filters,
+ partitionPredicates,
+ scanReportSupplier);
this.snapshot = snapshot;
this.branch = branch;
if (scan == null) {
@@ -140,6 +151,7 @@ class SparkCopyOnWriteScan extends
SparkPartitioningAwareScan<FileScanTask>
&& Objects.equals(branch, that.branch)
&& readSchema().equals(that.readSchema()) // compare Spark schemas to
ignore field ids
&& filtersDesc().equals(that.filtersDesc())
+ && partitionPredicates().equals(that.partitionPredicates())
&& Objects.equals(filteredLocations, that.filteredLocations);
}
@@ -152,6 +164,7 @@ class SparkCopyOnWriteScan extends
SparkPartitioningAwareScan<FileScanTask>
branch,
readSchema(),
filtersDesc(),
+ partitionPredicates(),
filteredLocations);
}
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkIncrementalAppendScan.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkIncrementalAppendScan.java
index ecfe782a53..932639ddc3 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkIncrementalAppendScan.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkIncrementalAppendScan.java
@@ -28,6 +28,7 @@ import org.apache.iceberg.expressions.Expression;
import org.apache.iceberg.metrics.ScanReport;
import org.apache.iceberg.spark.SparkReadConf;
import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
class SparkIncrementalAppendScan extends SparkRuntimeFilterableScan {
@@ -43,8 +44,18 @@ class SparkIncrementalAppendScan extends
SparkRuntimeFilterableScan {
SparkReadConf readConf,
Schema projection,
List<Expression> filters,
+ List<PartitionPredicate> partitionPredicates,
Supplier<ScanReport> scanReportSupplier) {
- super(spark, table, table.schema(), scan, readConf, projection, filters,
scanReportSupplier);
+ super(
+ spark,
+ table,
+ table.schema(),
+ scan,
+ readConf,
+ projection,
+ filters,
+ partitionPredicates,
+ scanReportSupplier);
this.startSnapshotId = startSnapshotId;
this.endSnapshotId = endSnapshotId;
}
@@ -66,6 +77,7 @@ class SparkIncrementalAppendScan extends
SparkRuntimeFilterableScan {
&& Objects.equals(endSnapshotId, that.endSnapshotId)
&& readSchema().equals(that.readSchema()) // compare Spark schemas to
ignore field ids
&& filtersDesc().equals(that.filtersDesc())
+ && partitionPredicates().equals(that.partitionPredicates())
&& runtimeFiltersDesc().equals(that.runtimeFiltersDesc());
}
@@ -78,6 +90,7 @@ class SparkIncrementalAppendScan extends
SparkRuntimeFilterableScan {
endSnapshotId,
readSchema(),
filtersDesc(),
+ partitionPredicates(),
runtimeFiltersDesc());
}
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkPartitioningAwareScan.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkPartitioningAwareScan.java
index fe5eeee8fb..124f00f0b9 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkPartitioningAwareScan.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkPartitioningAwareScan.java
@@ -34,6 +34,7 @@ import org.apache.iceberg.Scan;
import org.apache.iceberg.ScanTask;
import org.apache.iceberg.ScanTaskGroup;
import org.apache.iceberg.Schema;
+import org.apache.iceberg.StructLike;
import org.apache.iceberg.Table;
import org.apache.iceberg.exceptions.ValidationException;
import org.apache.iceberg.expressions.Expression;
@@ -46,9 +47,13 @@ import org.apache.iceberg.spark.SparkReadConf;
import org.apache.iceberg.types.Types.NestedField;
import org.apache.iceberg.types.Types.StructType;
import org.apache.iceberg.util.StructLikeSet;
+import org.apache.iceberg.util.StructProjection;
import org.apache.iceberg.util.TableScanUtil;
import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.PartitionFieldReference;
import org.apache.spark.sql.connector.expressions.Transform;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
import org.apache.spark.sql.connector.read.SupportsReportPartitioning;
import org.apache.spark.sql.connector.read.partitioning.KeyGroupedPartitioning;
import org.apache.spark.sql.connector.read.partitioning.Partitioning;
@@ -63,6 +68,9 @@ abstract class SparkPartitioningAwareScan<T extends
PartitionScanTask> extends S
private final Scan<?, ? extends ScanTask, ? extends ScanTaskGroup<?>> scan;
private final boolean preserveDataGrouping;
+ private final List<PartitionPredicate> partitionPredicates;
+ private final List<PartitionField> partitionPredicateFields;
+ private final StructType partitionPredicateType;
private Set<PartitionSpec> specs = null; // lazy cache of scanned specs
private List<T> tasks = null; // lazy cache of uncombined tasks
@@ -78,6 +86,7 @@ abstract class SparkPartitioningAwareScan<T extends
PartitionScanTask> extends S
SparkReadConf readConf,
Schema projection,
List<Expression> filters,
+ List<PartitionPredicate> partitionPredicates,
Supplier<ScanReport> scanReportSupplier) {
super(
spark,
@@ -91,6 +100,9 @@ abstract class SparkPartitioningAwareScan<T extends
PartitionScanTask> extends S
this.scan = scan;
this.preserveDataGrouping = readConf.preserveDataGrouping();
+ this.partitionPredicates = Lists.newArrayList(partitionPredicates);
+ this.partitionPredicateFields = activePartitionFields(table.spec());
+ this.partitionPredicateType = activePartitionType(table.spec());
if (scan == null) {
this.specs = Collections.emptySet();
@@ -181,6 +193,8 @@ abstract class SparkPartitioningAwareScan<T extends
PartitionScanTask> extends S
if (tasks == null) {
try (CloseableIterable<? extends ScanTask> taskIterable =
scan.planFiles()) {
List<T> plannedTasks = Lists.newArrayList();
+ Map<Integer, PartitionPredicateEvaluator> evaluatorsBySpecId =
Maps.newHashMap();
+ int numCandidateTasks = 0;
for (ScanTask task : taskIterable) {
ValidationException.check(
@@ -189,10 +203,25 @@ abstract class SparkPartitioningAwareScan<T extends
PartitionScanTask> extends S
taskJavaClass().getName(),
task.getClass().getName());
- plannedTasks.add(taskJavaClass().cast(task));
+ T partitionTask = taskJavaClass().cast(task);
+ numCandidateTasks += 1;
+
+ if (partitionPredicates.isEmpty()
+ || matchesPartitionPredicates(partitionTask,
evaluatorsBySpecId)) {
+ plannedTasks.add(partitionTask);
+ }
}
this.tasks = plannedTasks;
+
+ if (plannedTasks.size() < numCandidateTasks) {
+ LOG.info(
+ "{} of {} task(s) for table {} matched {} opaque Spark partition
predicate(s)",
+ plannedTasks.size(),
+ numCandidateTasks,
+ table().name(),
+ partitionPredicates.size());
+ }
} catch (IOException e) {
throw new UncheckedIOException("Failed to close scan: " + scan, e);
}
@@ -264,4 +293,98 @@ abstract class SparkPartitioningAwareScan<T extends
PartitionScanTask> extends S
.map(NestedField::name)
.collect(Collectors.joining(", "));
}
+
+ protected List<PartitionPredicate> partitionPredicates() {
+ return partitionPredicates;
+ }
+
+ private boolean matchesPartitionPredicates(
+ T task, Map<Integer, PartitionPredicateEvaluator> evaluatorsBySpecId) {
+ PartitionPredicateEvaluator evaluator =
+ evaluatorsBySpecId.computeIfAbsent(
+ task.spec().specId(), ignored -> new
PartitionPredicateEvaluator(task.spec()));
+ return evaluator.eval(task.partition());
+ }
+
+ private static List<PartitionField> activePartitionFields(PartitionSpec
spec) {
+ return spec.fields().stream()
+ .filter(field -> !field.transform().isVoid())
+ .collect(Collectors.toList());
+ }
+
+ private static StructType activePartitionType(PartitionSpec spec) {
+ List<NestedField> activeFields = Lists.newArrayList();
+ List<NestedField> partitionTypeFields = spec.partitionType().fields();
+
+ for (int pos = 0; pos < spec.fields().size(); pos += 1) {
+ if (!spec.fields().get(pos).transform().isVoid()) {
+ activeFields.add(partitionTypeFields.get(pos));
+ }
+ }
+
+ return StructType.of(activeFields);
+ }
+
+ private class PartitionPredicateEvaluator {
+ private final StructProjection partitionProjection;
+ private final StructInternalRow partitionKey;
+ private final boolean[] evaluablePredicates;
+
+ private PartitionPredicateEvaluator(PartitionSpec spec) {
+ this.partitionProjection =
+ StructProjection.createAllowMissing(spec.partitionType(),
partitionPredicateType);
+ this.partitionKey = new StructInternalRow(partitionPredicateType);
+ this.evaluablePredicates = findEvaluablePredicates(spec);
+ }
+
+ private boolean eval(StructLike partition) {
+ partitionKey.setStruct(partitionProjection.wrap(partition));
+
+ for (int pos = 0; pos < partitionPredicates.size(); pos += 1) {
+ if (evaluablePredicates[pos] &&
!partitionPredicates.get(pos).eval(partitionKey)) {
+ return false;
+ }
+ }
+
+ return true;
+ }
+
+ private boolean[] findEvaluablePredicates(PartitionSpec spec) {
+ Map<Integer, PartitionField> fieldsById = Maps.newHashMap();
+ for (PartitionField field : spec.fields()) {
+ if (!field.transform().isVoid()) {
+ fieldsById.put(field.fieldId(), field);
+ }
+ }
+
+ boolean[] result = new boolean[partitionPredicates.size()];
+ for (int pos = 0; pos < partitionPredicates.size(); pos += 1) {
+ result[pos] = canEvaluatePredicate(partitionPredicates.get(pos),
fieldsById);
+ }
+
+ return result;
+ }
+
+ private boolean canEvaluatePredicate(
+ PartitionPredicate predicate, Map<Integer, PartitionField> fieldsById)
{
+ for (NamedReference reference : predicate.references()) {
+ if (!(reference instanceof PartitionFieldReference
partitionReference)) {
+ return false;
+ }
+
+ int ordinal = partitionReference.ordinal();
+ if (ordinal < 0 || ordinal >= partitionPredicateFields.size()) {
+ return false;
+ }
+
+ PartitionField currentField = partitionPredicateFields.get(ordinal);
+ PartitionField specField = fieldsById.get(currentField.fieldId());
+ if (specField == null ||
!specField.transform().equals(currentField.transform())) {
+ return false;
+ }
+ }
+
+ return true;
+ }
+ }
}
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkRuntimeFilterableScan.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkRuntimeFilterableScan.java
index dfd38a1548..130e453196 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkRuntimeFilterableScan.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkRuntimeFilterableScan.java
@@ -53,6 +53,7 @@ import org.apache.iceberg.util.ContentFileUtil;
import org.apache.iceberg.util.DeleteFileSet;
import org.apache.spark.sql.SparkSession;
import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
import org.apache.spark.sql.connector.expressions.filter.Predicate;
import org.apache.spark.sql.connector.read.SupportsRuntimeV2Filtering;
import org.slf4j.Logger;
@@ -73,8 +74,18 @@ abstract class SparkRuntimeFilterableScan extends
SparkPartitioningAwareScan<Par
SparkReadConf readConf,
Schema projection,
List<Expression> filters,
+ List<PartitionPredicate> partitionPredicates,
Supplier<ScanReport> scanReportSupplier) {
- super(spark, table, schema, scan, readConf, projection, filters,
scanReportSupplier);
+ super(
+ spark,
+ table,
+ schema,
+ scan,
+ readConf,
+ projection,
+ filters,
+ partitionPredicates,
+ scanReportSupplier);
this.runtimeFilters = Lists.newArrayList();
}
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
index 6b9e314d35..d8e77ad7ff 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/SparkScanBuilder.java
@@ -125,6 +125,11 @@ public class SparkScanBuilder extends BaseSparkScanBuilder
SparkTableUtil.validateReadBranch(spark, table, branch, options);
}
+ @Override
+ public boolean supportsIterativePushdown() {
+ return true;
+ }
+
@Override
public boolean pushAggregation(Aggregation aggregation) {
if (!canPushDownAggregation(aggregation)) {
@@ -266,6 +271,7 @@ public class SparkScanBuilder extends BaseSparkScanBuilder
readConf(),
projection,
filters(),
+ partitionPredicates(),
metricsReporter()::scanReport);
}
@@ -280,6 +286,7 @@ public class SparkScanBuilder extends BaseSparkScanBuilder
readConf(),
projection,
filters(),
+ partitionPredicates(),
metricsReporter()::scanReport);
}
@@ -295,6 +302,7 @@ public class SparkScanBuilder extends BaseSparkScanBuilder
readConf(),
projection,
filters(),
+ partitionPredicates(),
metricsReporter()::scanReport);
}
diff --git
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/StructInternalRow.java
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/StructInternalRow.java
index 404cbc6519..78015cee1e 100644
---
a/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/StructInternalRow.java
+++
b/spark/v4.2/spark/src/main/java/org/apache/iceberg/spark/source/StructInternalRow.java
@@ -35,6 +35,7 @@ import org.apache.iceberg.StructLike;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.relocated.com.google.common.collect.Lists;
import org.apache.iceberg.spark.SparkSchemaUtil;
+import org.apache.iceberg.spark.SparkUtil;
import org.apache.iceberg.types.Type;
import org.apache.iceberg.types.Types;
import org.apache.iceberg.util.ByteBuffers;
@@ -183,8 +184,9 @@ class StructInternalRow extends InternalRow {
}
private UTF8String getUTF8StringInternal(int ordinal) {
- CharSequence seq = struct.get(ordinal, CharSequence.class);
- return UTF8String.fromString(seq.toString());
+ return (UTF8String)
+ SparkUtil.internalToSpark(
+ type.fields().get(ordinal).type(), struct.get(ordinal,
Object.class));
}
@Override
diff --git
a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestFilteredScan.java
b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestFilteredScan.java
index 24fecf4eb2..57b833a16c 100644
---
a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestFilteredScan.java
+++
b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/source/TestFilteredScan.java
@@ -61,7 +61,11 @@ import org.apache.spark.sql.Dataset;
import org.apache.spark.sql.Encoders;
import org.apache.spark.sql.Row;
import org.apache.spark.sql.SparkSession;
+import org.apache.spark.sql.catalyst.InternalRow;
import org.apache.spark.sql.catalyst.expressions.UnsafeRow;
+import org.apache.spark.sql.connector.expressions.NamedReference;
+import org.apache.spark.sql.connector.expressions.PartitionFieldReference;
+import org.apache.spark.sql.connector.expressions.filter.PartitionPredicate;
import org.apache.spark.sql.connector.expressions.filter.Predicate;
import org.apache.spark.sql.connector.read.Batch;
import org.apache.spark.sql.connector.read.InputPartition;
@@ -75,6 +79,7 @@ import org.apache.spark.sql.sources.LessThan;
import org.apache.spark.sql.sources.Not;
import org.apache.spark.sql.sources.StringStartsWith;
import org.apache.spark.sql.util.CaseInsensitiveStringMap;
+import org.apache.spark.unsafe.types.UTF8String;
import org.assertj.core.api.AbstractObjectAssert;
import org.junit.jupiter.api.AfterAll;
import org.junit.jupiter.api.BeforeAll;
@@ -196,10 +201,11 @@ public class TestFilteredScan {
public void testUnpartitionedIDFilters() {
CaseInsensitiveStringMap options =
new CaseInsensitiveStringMap(ImmutableMap.of("path",
unpartitioned.toString()));
- SparkScanBuilder builder =
- new SparkScanBuilder(spark, TABLES.load(options.get("path")), options);
for (int i = 0; i < 10; i += 1) {
+ SparkScanBuilder builder =
+ new SparkScanBuilder(spark, TABLES.load(options.get("path")),
options);
+
pushFilters(builder, EqualTo.apply("id", i));
Batch scan = builder.build().toBatch();
@@ -559,6 +565,53 @@ public class TestFilteredScan {
assertThat(scan.planInputPartitions()).hasSize(9);
}
+ @TestTemplate
+ public void partitionPredicatePrunesIdentityPartitions() {
+ Table table = buildPartitionedTable("partition_predicate",
PARTITION_BY_DATA);
+ CaseInsensitiveStringMap options =
+ new CaseInsensitiveStringMap(ImmutableMap.of("path",
table.location()));
+ SparkScanBuilder builder = new SparkScanBuilder(spark, table, options);
+ Predicate dataPredicate = GreaterThan.apply("id", -1L).toV2();
+ PartitionPredicate partitionPredicate = dataEndsWith("tion");
+
+ assertThat(builder.supportsIterativePushdown()).isTrue();
+ assertThat(builder.pushPredicates(new Predicate[] {dataPredicate}))
+ .containsExactly(dataPredicate);
+ assertThat(builder.pushPredicates(new Predicate[] {partitionPredicate}))
+ .containsExactly(partitionPredicate);
+ assertThat(builder.pushedPredicates()).containsExactly(dataPredicate,
partitionPredicate);
+
+ assertThat(builder.build().toBatch().planInputPartitions()).hasSize(1);
+ }
+
+ @TestTemplate
+ public void partitionPredicateKeepsTasksWithMissingFields() {
+ assumeThat(fileFormat)
+ .as("no need to run this across the entire test matrix")
+ .isEqualTo(FileFormat.PARQUET);
+
+ Table table = TABLES.load(unpartitioned.toString());
+ Dataset<Row> rows =
spark.read().format("iceberg").load(table.location()).cache();
+ rows.count();
+
+ table.updateSpec().addField("data").commit();
+ table.updateProperties().set("read.split.target-size", "2048").commit();
+ rows.write().format("iceberg").mode("append").save(table.location());
+ rows.unpersist();
+ table.refresh();
+
+ CaseInsensitiveStringMap options =
+ new CaseInsensitiveStringMap(ImmutableMap.of("path",
table.location()));
+ SparkScanBuilder builder = new SparkScanBuilder(spark, table, options);
+ PartitionPredicate partitionPredicate = dataEndsWith("tion");
+
+ assertThat(builder.pushPredicates(new Predicate[] {partitionPredicate}))
+ .containsExactly(partitionPredicate);
+
+ // The matching current-spec task and the old unpartitioned task must both
remain.
+ assertThat(builder.build().toBatch().planInputPartitions()).hasSize(2);
+ }
+
@TestTemplate
public void testPartitionedByIdStartsWith() {
Table table = buildPartitionedTable("partitioned_by_id", PARTITION_BY_ID);
@@ -671,6 +724,39 @@ public class TestFilteredScan {
filterable.pushPredicates(Arrays.stream(filters).map(Filter::toV2).toArray(Predicate[]::new));
}
+ private PartitionPredicate dataEndsWith(String suffix) {
+ PartitionFieldReference dataField =
+ new PartitionFieldReference() {
+ @Override
+ public int ordinal() {
+ return 0;
+ }
+
+ @Override
+ public String[] fieldNames() {
+ return new String[] {"data"};
+ }
+ };
+
+ return new PartitionPredicate() {
+ @Override
+ public String describe() {
+ return "data ends with " + suffix;
+ }
+
+ @Override
+ public NamedReference[] references() {
+ return new NamedReference[] {dataField};
+ }
+
+ @Override
+ public boolean eval(InternalRow partitionKey) {
+ UTF8String value = partitionKey.getUTF8String(0);
+ return value != null && value.toString().endsWith(suffix);
+ }
+ };
+ }
+
private Table buildPartitionedTable(String desc, PartitionSpec spec) {
File location = new File(parent, desc);
Table table = TABLES.create(SCHEMA, spec, location.toString());
diff --git
a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/sql/TestFilterPushDown.java
b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/sql/TestFilterPushDown.java
index 55299f7b67..55d5199c96 100644
---
a/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/sql/TestFilterPushDown.java
+++
b/spark/v4.2/spark/src/test/java/org/apache/iceberg/spark/sql/TestFilterPushDown.java
@@ -28,11 +28,15 @@ import java.nio.ByteOrder;
import java.sql.Timestamp;
import java.time.Instant;
import java.util.List;
+import java.util.Locale;
import org.apache.iceberg.Parameter;
import org.apache.iceberg.ParameterizedTestExtension;
import org.apache.iceberg.Parameters;
+import org.apache.iceberg.PartitionSpec;
import org.apache.iceberg.PlanningMode;
+import org.apache.iceberg.Schema;
import org.apache.iceberg.Table;
+import org.apache.iceberg.TableProperties;
import org.apache.iceberg.expressions.Expressions;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
import org.apache.iceberg.spark.SparkCatalogConfig;
@@ -43,10 +47,13 @@ import org.apache.iceberg.variants.ShreddedObject;
import org.apache.iceberg.variants.VariantMetadata;
import org.apache.iceberg.variants.Variants;
import org.apache.spark.sql.execution.SparkPlan;
+import org.apache.spark.sql.execution.datasources.v2.BatchScanExec;
+import org.apache.spark.sql.types.DataTypes;
import org.apache.spark.unsafe.types.VariantVal;
import org.junit.jupiter.api.AfterEach;
import org.junit.jupiter.api.TestTemplate;
import org.junit.jupiter.api.extension.ExtendWith;
+import scala.collection.JavaConverters;
@ExtendWith(ParameterizedTestExtension.class)
public class TestFilterPushDown extends TestBaseWithCatalog {
@@ -209,6 +216,178 @@ public class TestFilterPushDown extends
TestBaseWithCatalog {
ImmutableList.of(row(5, 500, "d5")));
}
+ @TestTemplate
+ public void partitionPredicatePushdownWithUDF() {
+ sql("CREATE TABLE %s (id INT, dep STRING) USING iceberg PARTITIONED BY
(dep)", tableName);
+ configurePlanningMode(planningMode);
+
+ setSmallSplitSize();
+ sql("INSERT INTO %s VALUES (1, 'd1')", tableName);
+ sql("INSERT INTO %s VALUES (2, 'd2')", tableName);
+ sql("INSERT INTO %s VALUES (3, 'd3')", tableName);
+
+ spark
+ .udf()
+ .register(
+ "is_d2", (String value) -> value != null && value.equals("d2"),
DataTypes.BooleanType);
+
+ SparkPlan plan =
+ executeAndKeepPlan(
+ () ->
+ assertEquals(
+ "Rows must match",
+ ImmutableList.of(row(2, "d2")),
+ sql("SELECT * FROM %s WHERE is_d2(dep)", tableName)));
+
+ assertInputPartitions(plan, 1);
+ }
+
+ @TestTemplate
+ public void partitionPredicatePushdownWithUnsupportedV2Filter() {
+ sql("CREATE TABLE %s (id INT, dep STRING) USING iceberg PARTITIONED BY
(dep)", tableName);
+ configurePlanningMode(planningMode);
+
+ setSmallSplitSize();
+ sql("INSERT INTO %s VALUES (1, 'd1')", tableName);
+ sql("INSERT INTO %s VALUES (2, 'd2')", tableName);
+ sql("INSERT INTO %s VALUES (3, 'd3')", tableName);
+
+ SparkPlan plan =
+ executeAndKeepPlan(
+ () ->
+ assertEquals(
+ "Rows must match",
+ ImmutableList.of(row(2, "d2")),
+ sql("SELECT * FROM %s WHERE dep LIKE '%%2'", tableName)));
+
+ assertInputPartitions(plan, 1);
+ }
+
+ @TestTemplate
+ void partitionPredicatePushdownForUUIDField() {
+ Schema schema = new Schema(Types.NestedField.optional(1, "uuid",
Types.UUIDType.get()));
+ PartitionSpec spec =
PartitionSpec.builderFor(schema).identity("uuid").build();
+ validationCatalog.createTable(tableIdent, schema, spec);
+ configurePlanningMode(planningMode);
+
+ setSmallSplitSize();
+ String matchingUuid = "00000000-0000-0000-0000-000000000001";
+ String otherUuid = "00000000-0000-0000-0000-000000000002";
+ sql("INSERT INTO %s VALUES ('%s')", tableName, matchingUuid);
+ sql("INSERT INTO %s VALUES ('%s')", tableName, otherUuid);
+
+ SparkPlan plan =
+ executeAndKeepPlan(
+ () ->
+ assertThat(sql("SELECT uuid FROM %s WHERE uuid LIKE '%%1'",
tableName))
+ .containsExactly(row(matchingUuid)));
+
+ assertInputPartitions(plan, 1);
+ }
+
+ @TestTemplate
+ public void partitionPredicatePushdownForNonFirstFields() {
+ sql(
+ "CREATE TABLE %s (p0 STRING, p1 STRING, p2 STRING, data STRING) "
+ + "USING iceberg PARTITIONED BY (p0, p1, p2)",
+ tableName);
+ configurePlanningMode(planningMode);
+
+ setSmallSplitSize();
+ sql("INSERT INTO %s VALUES ('a', 'x', '1', 'd1')", tableName);
+ sql("INSERT INTO %s VALUES ('a', 'y', '1', 'd2')", tableName);
+ sql("INSERT INTO %s VALUES ('a', 'x', '2', 'd3')", tableName);
+ sql("INSERT INTO %s VALUES ('b', 'x', '1', 'd4')", tableName);
+
+ SparkPlan nonFirstFieldPlan =
+ executeAndKeepPlan(
+ () ->
+ assertThat(sql("SELECT * FROM %s WHERE p1 LIKE '%%x'",
tableName))
+ .as("Rows must match")
+ .containsExactlyInAnyOrder(
+ row("a", "x", "1", "d1"),
+ row("a", "x", "2", "d3"),
+ row("b", "x", "1", "d4")));
+ assertInputPartitions(nonFirstFieldPlan, 3);
+
+ spark
+ .udf()
+ .register(
+ "concat_partition_fields",
+ (String left, String right) -> left == null || right == null ?
null : left + right,
+ DataTypes.StringType);
+
+ SparkPlan multipleFieldsPlan =
+ executeAndKeepPlan(
+ () ->
+ assertThat(
+ sql(
+ "SELECT * FROM %s WHERE
concat_partition_fields(p1, p2) = 'x1'",
+ tableName))
+ .as("Rows must match")
+ .containsExactlyInAnyOrder(row("a", "x", "1", "d1"),
row("b", "x", "1", "d4")));
+ assertInputPartitions(multipleFieldsPlan, 2);
+ }
+
+ @TestTemplate
+ public void partitionPredicatePushdownAfterDroppingMiddlePartitionField() {
+ sql(
+ "CREATE TABLE %s (p0 STRING, p1 STRING, p2 STRING, data STRING) "
+ + "USING iceberg PARTITIONED BY (p0, p1, p2)",
+ tableName);
+ configurePlanningMode(planningMode);
+
+ setSmallSplitSize();
+ sql("INSERT INTO %s VALUES ('a', 'x', 'keep', 'old-match')", tableName);
+ sql("INSERT INTO %s VALUES ('a', 'y', 'drop', 'old-non-match')",
tableName);
+
+ Table table = validationCatalog.loadTable(tableIdent);
+ table.updateSpec().removeField("p1").commit();
+ sql("REFRESH TABLE %s", tableName);
+
+ sql("INSERT INTO %s VALUES ('b', 'z', 'keep', 'new-match')", tableName);
+ sql("INSERT INTO %s VALUES ('b', 'z', 'drop', 'new-non-match')",
tableName);
+
+ SparkPlan plan =
+ executeAndKeepPlan(
+ () ->
+ assertThat(sql("SELECT * FROM %s WHERE p2 LIKE '%%keep'",
tableName))
+ .as("Rows must match")
+ .containsExactlyInAnyOrder(
+ row("a", "x", "keep", "old-match"), row("b", "z",
"keep", "new-match")));
+ assertInputPartitions(plan, 2);
+ }
+
+ @TestTemplate
+ public void partitionPredicatePushdownForNestedField() {
+ sql(
+ "CREATE TABLE %s (s STRUCT<tz: STRING, x: INT>, data STRING) "
+ + "USING iceberg PARTITIONED BY (s.tz)",
+ tableName);
+ configurePlanningMode(planningMode);
+
+ setSmallSplitSize();
+ sql("INSERT INTO %s VALUES (named_struct('tz', 'LA', 'x', 1), 'a')",
tableName);
+ sql("INSERT INTO %s VALUES (named_struct('tz', 'NY', 'x', 2), 'b')",
tableName);
+
+ spark
+ .udf()
+ .register(
+ "uppercase_tz",
+ (String value) -> value != null ? value.toUpperCase(Locale.ROOT) :
null,
+ DataTypes.StringType);
+
+ SparkPlan plan =
+ executeAndKeepPlan(
+ () ->
+ assertEquals(
+ "Rows must match",
+ ImmutableList.of(row(row("LA", 1), "a")),
+ sql("SELECT * FROM %s WHERE uppercase_tz(s.tz) = 'LA'",
tableName)));
+
+ assertInputPartitions(plan, 1);
+ }
+
@TestTemplate
public void testFilterPushdownWithHoursTransform() {
sql(
@@ -731,6 +910,18 @@ public class TestFilterPushDown extends
TestBaseWithCatalog {
checkFilters(predicate, null, icebergFilters, expectedRows);
}
+ private void setSmallSplitSize() {
+ Table table = validationCatalog.loadTable(tableIdent);
+ table.updateProperties().set(TableProperties.SPLIT_SIZE, "2048").commit();
+ }
+
+ private void assertInputPartitions(SparkPlan plan, int
expectedNumPartitions) {
+ List<SparkPlan> leaves =
JavaConverters.seqAsJavaListConverter(plan.collectLeaves()).asJava();
+ assertThat(leaves).singleElement().isInstanceOf(BatchScanExec.class);
+ BatchScanExec scan = (BatchScanExec) leaves.get(0);
+ assertThat(scan.inputPartitions().size()).isEqualTo(expectedNumPartitions);
+ }
+
private void checkFilters(
String predicate, String sparkFilter, String icebergFilters,
List<Object[]> expectedRows) {
