RussellSpitzer commented on code in PR #14948: URL: https://github.com/apache/iceberg/pull/14948#discussion_r3983793517
########## spark/v4.1/spark/src/test/java/org/apache/iceberg/spark/source/TestMergingSortedRowDataReader.java: ########## @@ -0,0 +1,802 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.spark.source; + +import static org.apache.iceberg.types.Types.NestedField.required; +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.when; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.util.Comparator; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.function.Function; +import java.util.stream.Stream; +import org.apache.iceberg.BaseScanTaskGroup; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.DataFiles; +import org.apache.iceberg.DeleteFile; +import org.apache.iceberg.FileScanTask; +import org.apache.iceberg.Files; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.ScanTaskGroup; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.Table; +import org.apache.iceberg.TableUtil; +import org.apache.iceberg.catalog.TableIdentifier; +import org.apache.iceberg.data.FileHelpers; +import org.apache.iceberg.data.GenericRecord; +import org.apache.iceberg.data.Record; +import org.apache.iceberg.expressions.Expressions; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.spark.TestBase; +import org.apache.iceberg.transforms.Transform; +import org.apache.iceberg.transforms.Transforms; +import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.Pair; +import org.apache.spark.rdd.InputFileBlockHolder; +import org.apache.spark.sql.catalyst.InternalRow; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.mockito.Mockito; + +class TestMergingSortedRowDataReader extends TestBase { + + private static final Schema SCHEMA = + new Schema( + required(1, "id", Types.IntegerType.get()), required(2, "data", Types.StringType.get())); + + private static final PartitionSpec SPEC = PartitionSpec.unpartitioned(); + + private Table table; + + @TempDir private Path temp; + + @BeforeEach + void before() { + table = catalog.createTable(TableIdentifier.of("default", "test_merging_reader"), SCHEMA, SPEC); + table.replaceSortOrder().asc("id").commit(); + } + + @AfterEach + void after() { + catalog.dropTable(TableIdentifier.of("default", "test_merging_reader")); + } + + @Test + void mergeTwoSortedFiles() throws IOException { + DataFile file1 = writeDataFile(record(1, "a"), record(3, "c"), record(5, "e")); + DataFile file2 = writeDataFile(record(2, "b"), record(4, "d"), record(6, "f")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(extractIds(rows)).containsExactly(1, 2, 3, 4, 5, 6); + } + + @Test + void mergeWithDuplicateKeys() throws IOException { + DataFile file1 = writeDataFile(record(1, "a"), record(2, "b")); + DataFile file2 = writeDataFile(record(1, "c"), record(2, "d")); + DataFile file3 = writeDataFile(record(1, "e"), record(3, "f")); + + table.newAppend().appendFile(file1).appendFile(file2).appendFile(file3).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(extractIds(rows)).containsExactly(1, 1, 1, 2, 2, 3); + } + + @Test + void mergeDescendingOrder() throws IOException { + catalog.dropTable(TableIdentifier.of("default", "test_merging_reader")); + table = catalog.createTable(TableIdentifier.of("default", "test_merging_reader"), SCHEMA, SPEC); + table.replaceSortOrder().desc("id").commit(); + + DataFile file1 = writeDataFile(record(6, "f"), record(4, "d")); + DataFile file2 = writeDataFile(record(5, "e"), record(3, "c"), record(1, "a")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(extractIds(rows)).containsExactly(6, 5, 4, 3, 1); + } + + @Test + void mergeWithNulls() throws IOException { + Schema nullableSchema = + new Schema( + Types.NestedField.optional(1, "id", Types.IntegerType.get()), + required(2, "data", Types.StringType.get())); + + catalog.dropTable(TableIdentifier.of("default", "test_merging_reader")); + table = + catalog.createTable( + TableIdentifier.of("default", "test_merging_reader"), nullableSchema, SPEC); + table.replaceSortOrder().asc("id").commit(); + + DataFile file1 = writeDataFile(nullRecord("x"), record(3, "c")); + DataFile file2 = writeDataFile(nullRecord("y"), record(1, "a"), record(2, "b")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(rows).hasSize(5); + assertThat(rows.get(0).isNullAt(0)).isTrue(); + assertThat(rows.get(1).isNullAt(0)).isTrue(); + assertThat(extractIds(rows.subList(2, 5))).containsExactly(1, 2, 3); + } + + @Test + void mergeThreeFiles() throws IOException { + DataFile file1 = writeDataFile(record(1, "a"), record(4, "d"), record(7, "g")); + DataFile file2 = writeDataFile(record(2, "b"), record(5, "e"), record(8, "h")); + DataFile file3 = writeDataFile(record(3, "c"), record(6, "f"), record(9, "i")); + + table.newAppend().appendFile(file1).appendFile(file2).appendFile(file3).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(extractIds(rows)).containsExactly(1, 2, 3, 4, 5, 6, 7, 8, 9); + } + + @Test + void mergeWithSortKeyNotInProjection() throws IOException { + DataFile file1 = writeDataFile(record(1, "a"), record(3, "c"), record(5, "e")); + DataFile file2 = writeDataFile(record(2, "b"), record(4, "d"), record(6, "f")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + // Project only "data". The sort key "id" is missing from the projection, so it is added to + // the read schema for the merge comparator and stripped from the rows returned to Spark. + Schema dataOnly = table.schema().select("data"); + List<InternalRow> rows = readMerged(table, dataOnly); + + // Rows come back ordered by id even though id is not projected. + assertThat(extractData(rows, 0)).containsExactly("a", "b", "c", "d", "e", "f"); + // Only the projected column is present in the returned rows. + assertThat(rows.get(0).numFields()).isEqualTo(1); + } + + @Test + void mergeAfterSortOrderEvolution() throws IOException { + // Evolve the sort order from "id" to "data". The reader should merge by the current order. + table.replaceSortOrder().asc("data").commit(); + + DataFile file1 = writeDataFile(record(5, "a"), record(3, "c"), record(1, "e")); + DataFile file2 = writeDataFile(record(6, "b"), record(4, "d"), record(2, "f")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + List<InternalRow> rows = readMerged(table); + + // Ordered by data, not by id. + assertThat(extractData(rows, 1)).containsExactly("a", "b", "c", "d", "e", "f"); + } + + @Test + void mergeWithStructColumnNotInSortOrder() throws IOException { + catalog.dropTable(TableIdentifier.of("default", "test_merging_reader")); + + Schema schemaWithStruct = + new Schema( + required(1, "id", Types.IntegerType.get()), + required(2, "data", Types.StringType.get()), + required( + 4, "location", Types.StructType.of(required(5, "city", Types.StringType.get())))); + + table = + catalog.createTable( + TableIdentifier.of("default", "test_merging_reader"), schemaWithStruct, SPEC); + table.replaceSortOrder().asc("id").commit(); + + DataFile file1 = writeDataFile(structRecord(1, "a", "NYC"), structRecord(3, "c", "SFO")); + DataFile file2 = writeDataFile(structRecord(2, "b", "LAX"), structRecord(4, "d", "SEA")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + // Project the struct but not the sort key, so the merge schema is widened around a struct. + Schema projection = table.schema().select("location"); + List<InternalRow> rows = readMerged(table, projection); + + assertThat(rows.get(0).numFields()).isEqualTo(1); + assertThat(rows.stream().map(row -> row.getStruct(0, 1).getUTF8String(0).toString()).toList()) + .containsExactly("NYC", "LAX", "SFO", "SEA"); + } + + @Test + void mergeWithArrayOfStructsDoesNotCorruptElements() throws IOException { + catalog.dropTable(TableIdentifier.of("default", "test_merging_reader")); + + Types.StructType element = Types.StructType.of(required(4, "a", Types.IntegerType.get())); + Schema arrayOfStructs = + new Schema( + required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional(2, "arr", Types.ListType.ofOptional(3, element))); + + table = + catalog.createTable( + TableIdentifier.of("default", "test_merging_reader"), arrayOfStructs, SPEC); + table.replaceSortOrder().asc("id").commit(); + + // File1 = [(1,[a=10]), (3,[a=30])], File2 = [(2,[a=20])]. SortedMerge advances file1's reader + // before returning the row for id=1, so a shallow copy would let id=3's struct clobber id=1's. + DataFile file1 = + writeDataFile( + arrayOfStructsRecord(arrayOfStructs, element, 1, 10), + arrayOfStructsRecord(arrayOfStructs, element, 3, 30)); + DataFile file2 = writeDataFile(arrayOfStructsRecord(arrayOfStructs, element, 2, 20)); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(extractIds(rows)).containsExactly(1, 2, 3); + assertThat(rows.stream().map(row -> row.getArray(1).getStruct(0, 1).getInt(0)).toList()) + .containsExactly(10, 20, 30); + } + + @Test + void mergeWithMapOfStructsDoesNotCorruptElements() throws IOException { + catalog.dropTable(TableIdentifier.of("default", "test_merging_reader")); + + Types.StructType element = Types.StructType.of(required(5, "a", Types.IntegerType.get())); + Schema mapOfStructs = + new Schema( + required(1, "id", Types.IntegerType.get()), + Types.NestedField.optional( + 2, "m", Types.MapType.ofOptional(3, 4, Types.StringType.get(), element))); + + table = + catalog.createTable( + TableIdentifier.of("default", "test_merging_reader"), mapOfStructs, SPEC); + table.replaceSortOrder().asc("id").commit(); + + DataFile file1 = + writeDataFile( + mapOfStructsRecord(mapOfStructs, element, 1, 10), + mapOfStructsRecord(mapOfStructs, element, 3, 30)); + DataFile file2 = writeDataFile(mapOfStructsRecord(mapOfStructs, element, 2, 20)); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + List<InternalRow> rows = readMerged(table); + + assertThat(extractIds(rows)).containsExactly(1, 2, 3); + assertThat( + rows.stream().map(row -> row.getMap(1).valueArray().getStruct(0, 1).getInt(0)).toList()) + .containsExactly(10, 20, 30); + } + + @Test + void mergeRejectsStaleSortOrderId() throws IOException { + SortOrder oldSortOrder = table.sortOrder(); + + // file1 keeps the old order id, file2 is written with the evolved one + DataFile file1 = + DataFiles.builder(table.spec()) + .copy(writeRecords(record(1, "a"), record(3, "c"))) + .withSortOrder(oldSortOrder) + .build(); + + table.replaceSortOrder().asc("data").commit(); + DataFile file2 = writeDataFile(record(2, "b"), record(4, "d")); + + table.newAppend().appendFile(file1).appendFile(file2).commit(); + + assertThatThrownBy(() -> readMerged(table)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Not all files in task group have the expected sort order"); + } + + @Test + void mergeRejectsMissingSortOrderId() { + // sort_order_id is optional in the manifest schema, so a file may report null + ScanTaskGroup<FileScanTask> taskGroup = + taskGroupWithSortOrderIds(table.sortOrder().orderId(), null); + + assertThatThrownBy( + () -> + new MergingSortedRowDataReader( + table, table.io(), taskGroup, table.schema(), true, false)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Not all files in task group have the expected sort order"); + } + + @Test + void mergeRejectsSingleFile() throws IOException { + DataFile file1 = writeDataFile(record(1, "a"), record(3, "c")); + + table.newAppend().appendFile(file1).commit(); + table.refresh(); + + BaseScanTaskGroup<FileScanTask> taskGroup = new BaseScanTaskGroup<>(planFiles(table)); + + assertThatThrownBy( + () -> + new MergingSortedRowDataReader( + table, table.io(), taskGroup, table.schema(), true, false)) + .isInstanceOf(IllegalArgumentException.class) + .hasMessageContaining("Merging reader requires multiple files, got 1"); + } + + @Test + void mergeRejectsUnsortedTable() throws IOException { Review Comment: Probably drop this since I don't think we should gate on table sort order, instead just check file sort order -- 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]
