Skip to content

Commit bf17005

Browse files
isha97isha97
committed
fix no partition test
1 parent 356c1ad commit bf17005

File tree

1 file changed

+25
-22
lines changed

1 file changed

+25
-22
lines changed

‎spark-bigquery-connector-common/src/test/java/com/google/cloud/spark/bigquery/integration/WriteIntegrationTestBase.java‎

Lines changed: 25 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -1579,7 +1579,8 @@ public void testWriteToTimestampField() {
15791579
assertThat(head.get(head.fieldIndex("timestamp1"))).isEqualTo(timestamp1);
15801580
}
15811581

1582-
protected Dataset<Row> writeAndLoadDatasetOverwriteDynamicPartition(Dataset<Row> df) {
1582+
protected Dataset<Row> writeAndLoadDatasetOverwriteDynamicPartition(
1583+
Dataset<Row> df, boolean isPartitioned) {
15831584
df.write()
15841585
.format("bigquery")
15851586
.mode(SaveMode.Overwrite)
@@ -1591,10 +1592,12 @@ protected Dataset<Row> writeAndLoadDatasetOverwriteDynamicPartition(Dataset<Row>
15911592
.option("temporaryGcsBucket", TestConstants.TEMPORARY_GCS_BUCKET)
15921593
.save();
15931594

1594-
IntegrationTestUtils.runQuery(
1595-
String.format(
1596-
"ALTER TABLE %s.%s SET OPTIONS (require_partition_filter = false)",
1597-
testDataset, testTable));
1595+
if (isPartitioned) {
1596+
IntegrationTestUtils.runQuery(
1597+
String.format(
1598+
"ALTER TABLE %s.%s SET OPTIONS (require_partition_filter = false)",
1599+
testDataset, testTable));
1600+
}
15981601

15991602
return spark
16001603
.read()
@@ -1626,7 +1629,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByHour() {
16261629
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
16271630
StructField.apply(orderDateTime, DataTypes.TimestampType, true, Metadata.empty())));
16281631

1629-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1632+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
16301633
assertThat(result.count()).isEqualTo(3);
16311634
List<Row> rows = result.collectAsList();
16321635
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1670,7 +1673,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByDay() {
16701673
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
16711674
StructField.apply(orderDateTime, DataTypes.TimestampType, true, Metadata.empty())));
16721675

1673-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1676+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
16741677
assertThat(result.count()).isEqualTo(3);
16751678
List<Row> rows = result.collectAsList();
16761679
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1714,7 +1717,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByMonth() {
17141717
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
17151718
StructField.apply(orderDateTime, DataTypes.TimestampType, true, Metadata.empty())));
17161719

1717-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1720+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
17181721
assertThat(result.count()).isEqualTo(3);
17191722
List<Row> rows = result.collectAsList();
17201723
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1758,7 +1761,7 @@ public void testOverwriteDynamicPartition_partitionTimestampByYear() {
17581761
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
17591762
StructField.apply(orderDateTime, DataTypes.TimestampType, true, Metadata.empty())));
17601763

1761-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1764+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
17621765
assertThat(result.count()).isEqualTo(3);
17631766
List<Row> rows = result.collectAsList();
17641767
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1801,7 +1804,7 @@ public void testOverwriteDynamicPartition_partitionDateByDay() {
18011804
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
18021805
StructField.apply(orderDate, DataTypes.DateType, true, Metadata.empty())));
18031806

1804-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1807+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
18051808
assertThat(result.count()).isEqualTo(3);
18061809
List<Row> rows = result.collectAsList();
18071810
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1842,7 +1845,7 @@ public void testOverwriteDynamicPartition_partitionDateByMonth() {
18421845
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
18431846
StructField.apply(orderDate, DataTypes.DateType, true, Metadata.empty())));
18441847

1845-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1848+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
18461849
assertThat(result.count()).isEqualTo(3);
18471850
List<Row> rows = result.collectAsList();
18481851
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1883,7 +1886,7 @@ public void testOverwriteDynamicPartition_partitionDateByYear() {
18831886
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
18841887
StructField.apply(orderDate, DataTypes.DateType, true, Metadata.empty())));
18851888

1886-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1889+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
18871890
assertThat(result.count()).isEqualTo(3);
18881891
List<Row> rows = result.collectAsList();
18891892
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1925,7 +1928,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByHour() {
19251928
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
19261929
StructField.apply(orderDateTime, timeStampNTZType.get(), true, Metadata.empty())));
19271930

1928-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1931+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
19291932
assertThat(result.count()).isEqualTo(3);
19301933
List<Row> rows = result.collectAsList();
19311934
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -1967,7 +1970,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByDay() {
19671970
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
19681971
StructField.apply(orderDateTime, timeStampNTZType.get(), true, Metadata.empty())));
19691972

1970-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
1973+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
19711974
assertThat(result.count()).isEqualTo(3);
19721975
List<Row> rows = result.collectAsList();
19731976
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2009,7 +2012,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByMonth() {
20092012
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
20102013
StructField.apply(orderDateTime, timeStampNTZType.get(), true, Metadata.empty())));
20112014

2012-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2015+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
20132016
assertThat(result.count()).isEqualTo(3);
20142017
List<Row> rows = result.collectAsList();
20152018
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2051,7 +2054,7 @@ public void testOverwriteDynamicPartition_partitionDateTimeByYear() {
20512054
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
20522055
StructField.apply(orderDateTime, timeStampNTZType.get(), true, Metadata.empty())));
20532056

2054-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2057+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
20552058
assertThat(result.count()).isEqualTo(3);
20562059
List<Row> rows = result.collectAsList();
20572060
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2091,7 +2094,7 @@ public void testOverwriteDynamicPartition_noTimePartitioning() {
20912094
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
20922095
StructField.apply(orderDateTime, DataTypes.TimestampType, true, Metadata.empty())));
20932096

2094-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2097+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, false);
20952098
assertThat(result.count()).isEqualTo(2);
20962099
List<Row> rows = result.collectAsList();
20972100
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2132,7 +2135,7 @@ public void testOverwriteDynamicPartition_rangePartitioned() {
21322135
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
21332136
StructField.apply(orderCount, DataTypes.IntegerType, true, Metadata.empty())));
21342137

2135-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2138+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
21362139
assertThat(result.count()).isEqualTo(5);
21372140
List<Row> rows = result.collectAsList();
21382141
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2179,7 +2182,7 @@ public void testOverwriteDynamicPartition_rangePartitionedOutsideRangeLessThanSt
21792182
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
21802183
StructField.apply(orderCount, DataTypes.IntegerType, true, Metadata.empty())));
21812184

2182-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2185+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
21832186
assertThat(result.count()).isEqualTo(2);
21842187
List<Row> rows = result.collectAsList();
21852188
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2214,7 +2217,7 @@ public void testOverwriteDynamicPartition_rangePartitionedOutsideRangeGreaterTha
22142217
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
22152218
StructField.apply(orderCount, DataTypes.IntegerType, true, Metadata.empty())));
22162219

2217-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2220+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
22182221
assertThat(result.count()).isEqualTo(2);
22192222
List<Row> rows = result.collectAsList();
22202223
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2249,7 +2252,7 @@ public void testOverwriteDynamicPartition_rangePartitionedBoundaryCondition() {
22492252
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
22502253
StructField.apply(orderCount, DataTypes.IntegerType, true, Metadata.empty())));
22512254

2252-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2255+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
22532256
assertThat(result.count()).isEqualTo(3);
22542257
List<Row> rows = result.collectAsList();
22552258
rows.sort(Comparator.comparing(row -> row.getLong(row.fieldIndex(orderId))));
@@ -2288,7 +2291,7 @@ public void testOverwriteDynamicPartition_rangePartitionedWithNulls() {
22882291
StructField.apply(orderId, DataTypes.IntegerType, true, Metadata.empty()),
22892292
StructField.apply(orderCount, DataTypes.IntegerType, true, Metadata.empty())));
22902293

2291-
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df);
2294+
Dataset<Row> result = writeAndLoadDatasetOverwriteDynamicPartition(df, true);
22922295
assertThat(result.count()).isEqualTo(3);
22932296

22942297
List<Row> rows = result.collectAsList();

0 commit comments

Comments
 (0)