Skip to content

Commit f2875fd

Browse files
authored
Spark: Backport split-size session conf to 3.5 and 4.0 (#17199)
1 parent 8b26043 commit f2875fd

8 files changed

Lines changed: 176 additions & 2 deletions

File tree

spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,13 +192,18 @@ public int orcBatchSize() {
192192
}
193193

194194
public Long splitSizeOption() {
195-
return confParser.longConf().option(SparkReadOptions.SPLIT_SIZE).parseOptional();
195+
return confParser
196+
.longConf()
197+
.option(SparkReadOptions.SPLIT_SIZE)
198+
.sessionConf(SparkSQLProperties.READ_SPLIT_SIZE)
199+
.parseOptional();
196200
}
197201

198202
public long splitSize() {
199203
return confParser
200204
.longConf()
201205
.option(SparkReadOptions.SPLIT_SIZE)
206+
.sessionConf(SparkSQLProperties.READ_SPLIT_SIZE)
202207
.tableProperty(TableProperties.SPLIT_SIZE)
203208
.defaultValue(TableProperties.SPLIT_SIZE_DEFAULT)
204209
.parse();

spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -113,6 +113,9 @@ private SparkSQLProperties() {}
113113
public static final String READ_ADAPTIVE_SPLIT_SIZE_ENABLED =
114114
"spark.sql.iceberg.read.adaptive-split-size.enabled";
115115

116+
// Overrides the split target size for scan planning
117+
public static final String READ_SPLIT_SIZE = "spark.sql.iceberg.read.split-size";
118+
116119
// Overrides the parallelism used for adaptive split sizing. When unset, the parallelism
117120
// defaults to max(spark.default.parallelism, spark.sql.shuffle.partitions).
118121
public static final String READ_ADAPTIVE_SPLIT_SIZE_PARALLELISM =

spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/TestSparkReadConf.java

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,44 @@ public void testSplitParallelismSessionConf() {
6767
});
6868
}
6969

70+
@TestTemplate
71+
public void testSplitSizeSessionConf() {
72+
Table table = validationCatalog.loadTable(tableIdent);
73+
withSQLConf(
74+
ImmutableMap.of(SparkSQLProperties.READ_SPLIT_SIZE, "42"),
75+
() -> {
76+
SparkReadConf conf = new SparkReadConf(spark, table, CaseInsensitiveStringMap.empty());
77+
assertThat(conf.splitSizeOption()).isEqualTo(42L);
78+
assertThat(conf.splitSize()).isEqualTo(42L);
79+
});
80+
}
81+
82+
@TestTemplate
83+
public void testSplitSizeCamelCaseSessionConf() {
84+
Table table = validationCatalog.loadTable(tableIdent);
85+
withSQLConf(
86+
ImmutableMap.of("spark.sql.iceberg.read.splitSize", "42"),
87+
() -> {
88+
SparkReadConf conf = new SparkReadConf(spark, table, CaseInsensitiveStringMap.empty());
89+
assertThat(conf.splitSizeOption()).isEqualTo(42L);
90+
assertThat(conf.splitSize()).isEqualTo(42L);
91+
});
92+
}
93+
94+
@TestTemplate
95+
public void testSplitSizeReadOptionOverridesSessionConf() {
96+
Table table = validationCatalog.loadTable(tableIdent);
97+
withSQLConf(
98+
ImmutableMap.of(SparkSQLProperties.READ_SPLIT_SIZE, "24"),
99+
() -> {
100+
CaseInsensitiveStringMap options =
101+
new CaseInsensitiveStringMap(ImmutableMap.of(SparkReadOptions.SPLIT_SIZE, "42"));
102+
SparkReadConf conf = new SparkReadConf(spark, table, options);
103+
assertThat(conf.splitSizeOption()).isEqualTo(42L);
104+
assertThat(conf.splitSize()).isEqualTo(42L);
105+
});
106+
}
107+
70108
@TestTemplate
71109
public void testSplitParallelismRejectsZero() {
72110
Table table = validationCatalog.loadTable(tableIdent);

spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkScan.java

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.apache.iceberg.Parameter;
3737
import org.apache.iceberg.ParameterizedTestExtension;
3838
import org.apache.iceberg.Parameters;
39+
import org.apache.iceberg.ScanTask;
3940
import org.apache.iceberg.Table;
4041
import org.apache.iceberg.TableProperties;
4142
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
@@ -144,6 +145,37 @@ public void testEstimatedRowCount() throws NoSuchTableException {
144145
assertThat(stats.numRows().getAsLong()).isEqualTo(10000L);
145146
}
146147

148+
@TestTemplate
149+
public void testSessionSplitSizeSkipsAdaptiveSplitSizeAdjustment() {
150+
sql("CREATE TABLE %s (id BIGINT) USING iceberg", tableName);
151+
152+
Table table = validationCatalog.loadTable(tableIdent);
153+
long splitSize = TableProperties.SPLIT_SIZE_DEFAULT;
154+
List<ScanTask> tasks = ImmutableList.of(scanTaskWithSize(16L * 1024 * 1024));
155+
156+
withSQLConf(
157+
ImmutableMap.of(SQLConf.SHUFFLE_PARTITIONS().key(), "200"),
158+
() -> {
159+
SparkScan scan =
160+
(SparkScan)
161+
new SparkScanBuilder(spark, table, CaseInsensitiveStringMap.empty()).build();
162+
assertThat(scan.adjustSplitSize(tasks, splitSize)).isLessThan(splitSize);
163+
});
164+
165+
withSQLConf(
166+
ImmutableMap.of(
167+
SQLConf.SHUFFLE_PARTITIONS().key(),
168+
"200",
169+
SparkSQLProperties.READ_SPLIT_SIZE,
170+
String.valueOf(splitSize)),
171+
() -> {
172+
SparkScan scan =
173+
(SparkScan)
174+
new SparkScanBuilder(spark, table, CaseInsensitiveStringMap.empty()).build();
175+
assertThat(scan.adjustSplitSize(tasks, splitSize)).isEqualTo(splitSize);
176+
});
177+
}
178+
147179
@TestTemplate
148180
public void testTableWithoutColStats() throws NoSuchTableException {
149181
sql("CREATE TABLE %s (id int, data string) USING iceberg", tableName);
@@ -1022,6 +1054,15 @@ public void testPartitionedOr() throws Exception {
10221054
assertThat(scan.planInputPartitions()).hasSize(4);
10231055
}
10241056

1057+
private ScanTask scanTaskWithSize(long sizeBytes) {
1058+
return new ScanTask() {
1059+
@Override
1060+
public long sizeBytes() {
1061+
return sizeBytes;
1062+
}
1063+
};
1064+
}
1065+
10251066
private SparkScanBuilder scanBuilder() throws Exception {
10261067
Table table = Spark3Util.loadIcebergTable(spark, tableName);
10271068
CaseInsensitiveStringMap options =

spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/SparkReadConf.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,13 +192,18 @@ public int orcBatchSize() {
192192
}
193193

194194
public Long splitSizeOption() {
195-
return confParser.longConf().option(SparkReadOptions.SPLIT_SIZE).parseOptional();
195+
return confParser
196+
.longConf()
197+
.option(SparkReadOptions.SPLIT_SIZE)
198+
.sessionConf(SparkSQLProperties.READ_SPLIT_SIZE)
199+
.parseOptional();
196200
}
197201

198202
public long splitSize() {
199203
return confParser
200204
.longConf()
201205
.option(SparkReadOptions.SPLIT_SIZE)
206+
.sessionConf(SparkSQLProperties.READ_SPLIT_SIZE)
202207
.tableProperty(TableProperties.SPLIT_SIZE)
203208
.defaultValue(TableProperties.SPLIT_SIZE_DEFAULT)
204209
.parse();

spark/v4.0/spark/src/main/java/org/apache/iceberg/spark/SparkSQLProperties.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -116,6 +116,9 @@ private SparkSQLProperties() {}
116116
public static final String READ_ADAPTIVE_SPLIT_SIZE_ENABLED =
117117
"spark.sql.iceberg.read.adaptive-split-size.enabled";
118118

119+
// Overrides the split target size for scan planning
120+
public static final String READ_SPLIT_SIZE = "spark.sql.iceberg.read.split-size";
121+
119122
// Overrides the parallelism used for adaptive split sizing. When unset, the parallelism
120123
// defaults to max(spark.default.parallelism, spark.sql.shuffle.partitions).
121124
public static final String READ_ADAPTIVE_SPLIT_SIZE_PARALLELISM =

spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/TestSparkReadConf.java

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,44 @@ public void testSplitParallelismSessionConf() {
6767
});
6868
}
6969

70+
@TestTemplate
71+
public void testSplitSizeSessionConf() {
72+
Table table = validationCatalog.loadTable(tableIdent);
73+
withSQLConf(
74+
ImmutableMap.of(SparkSQLProperties.READ_SPLIT_SIZE, "42"),
75+
() -> {
76+
SparkReadConf conf = new SparkReadConf(spark, table, CaseInsensitiveStringMap.empty());
77+
assertThat(conf.splitSizeOption()).isEqualTo(42L);
78+
assertThat(conf.splitSize()).isEqualTo(42L);
79+
});
80+
}
81+
82+
@TestTemplate
83+
public void testSplitSizeCamelCaseSessionConf() {
84+
Table table = validationCatalog.loadTable(tableIdent);
85+
withSQLConf(
86+
ImmutableMap.of("spark.sql.iceberg.read.splitSize", "42"),
87+
() -> {
88+
SparkReadConf conf = new SparkReadConf(spark, table, CaseInsensitiveStringMap.empty());
89+
assertThat(conf.splitSizeOption()).isEqualTo(42L);
90+
assertThat(conf.splitSize()).isEqualTo(42L);
91+
});
92+
}
93+
94+
@TestTemplate
95+
public void testSplitSizeReadOptionOverridesSessionConf() {
96+
Table table = validationCatalog.loadTable(tableIdent);
97+
withSQLConf(
98+
ImmutableMap.of(SparkSQLProperties.READ_SPLIT_SIZE, "24"),
99+
() -> {
100+
CaseInsensitiveStringMap options =
101+
new CaseInsensitiveStringMap(ImmutableMap.of(SparkReadOptions.SPLIT_SIZE, "42"));
102+
SparkReadConf conf = new SparkReadConf(spark, table, options);
103+
assertThat(conf.splitSizeOption()).isEqualTo(42L);
104+
assertThat(conf.splitSize()).isEqualTo(42L);
105+
});
106+
}
107+
70108
@TestTemplate
71109
public void testSplitParallelismRejectsZero() {
72110
Table table = validationCatalog.loadTable(tableIdent);

spark/v4.0/spark/src/test/java/org/apache/iceberg/spark/source/TestSparkScan.java

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@
3636
import org.apache.iceberg.Parameter;
3737
import org.apache.iceberg.ParameterizedTestExtension;
3838
import org.apache.iceberg.Parameters;
39+
import org.apache.iceberg.ScanTask;
3940
import org.apache.iceberg.Table;
4041
import org.apache.iceberg.TableProperties;
4142
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
@@ -144,6 +145,37 @@ public void testEstimatedRowCount() throws NoSuchTableException {
144145
assertThat(stats.numRows().getAsLong()).isEqualTo(10000L);
145146
}
146147

148+
@TestTemplate
149+
public void testSessionSplitSizeSkipsAdaptiveSplitSizeAdjustment() {
150+
sql("CREATE TABLE %s (id BIGINT) USING iceberg", tableName);
151+
152+
Table table = validationCatalog.loadTable(tableIdent);
153+
long splitSize = TableProperties.SPLIT_SIZE_DEFAULT;
154+
List<ScanTask> tasks = ImmutableList.of(scanTaskWithSize(16L * 1024 * 1024));
155+
156+
withSQLConf(
157+
ImmutableMap.of(SQLConf.SHUFFLE_PARTITIONS().key(), "200"),
158+
() -> {
159+
SparkScan scan =
160+
(SparkScan)
161+
new SparkScanBuilder(spark, table, CaseInsensitiveStringMap.empty()).build();
162+
assertThat(scan.adjustSplitSize(tasks, splitSize)).isLessThan(splitSize);
163+
});
164+
165+
withSQLConf(
166+
ImmutableMap.of(
167+
SQLConf.SHUFFLE_PARTITIONS().key(),
168+
"200",
169+
SparkSQLProperties.READ_SPLIT_SIZE,
170+
String.valueOf(splitSize)),
171+
() -> {
172+
SparkScan scan =
173+
(SparkScan)
174+
new SparkScanBuilder(spark, table, CaseInsensitiveStringMap.empty()).build();
175+
assertThat(scan.adjustSplitSize(tasks, splitSize)).isEqualTo(splitSize);
176+
});
177+
}
178+
147179
@TestTemplate
148180
public void testTableWithoutColStats() throws NoSuchTableException {
149181
sql("CREATE TABLE %s (id int, data string) USING iceberg", tableName);
@@ -1022,6 +1054,15 @@ public void testPartitionedOr() throws Exception {
10221054
assertThat(scan.planInputPartitions()).hasSize(4);
10231055
}
10241056

1057+
private ScanTask scanTaskWithSize(long sizeBytes) {
1058+
return new ScanTask() {
1059+
@Override
1060+
public long sizeBytes() {
1061+
return sizeBytes;
1062+
}
1063+
};
1064+
}
1065+
10251066
private SparkScanBuilder scanBuilder() throws Exception {
10261067
Table table = Spark3Util.loadIcebergTable(spark, tableName);
10271068
CaseInsensitiveStringMap options =

0 commit comments

Comments
 (0)