From 8d279fc15b5e507baa1ea2629911d486f4112613 Mon Sep 17 00:00:00 2001 From: ajantha-bhat Date: Mon, 27 Nov 2023 15:25:09 +0530 Subject: [PATCH 1/2] Core: Add a util to read write partition stats --- .../org/apache/iceberg/PartitionEntry.java | 361 ++++++++++++++++++ .../java/org/apache/iceberg/Partitioning.java | 5 +- .../apache/iceberg/TestPartitionEntry.java | 76 ++++ .../iceberg/data/PartitionStatsUtil.java | 148 +++++++ .../iceberg/data/TestPartitionStatsUtil.java | 125 ++++++ 5 files changed, 714 insertions(+), 1 deletion(-) create mode 100644 core/src/main/java/org/apache/iceberg/PartitionEntry.java create mode 100644 core/src/test/java/org/apache/iceberg/TestPartitionEntry.java create mode 100644 data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java create mode 100644 data/src/test/java/org/apache/iceberg/data/TestPartitionStatsUtil.java diff --git a/core/src/main/java/org/apache/iceberg/PartitionEntry.java b/core/src/main/java/org/apache/iceberg/PartitionEntry.java new file mode 100644 index 000000000000..ac6da58c5d5a --- /dev/null +++ b/core/src/main/java/org/apache/iceberg/PartitionEntry.java @@ -0,0 +1,361 @@ +/* + * 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; + +import java.util.Objects; +import org.apache.avro.Schema; +import org.apache.avro.generic.IndexedRecord; +import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.types.Types; + +public class PartitionEntry implements IndexedRecord { + private PartitionData partitionData; + private int specId; + private long dataRecordCount; + private int dataFileCount; + private long dataFileSizeInBytes; + private long posDeleteRecordCount; + private int posDeleteFileCount; + private long eqDeleteRecordCount; + private int eqDeleteFileCount; + // Optional accurate count of records in a partition after applying the delete files if any + private long totalRecordCount; + // Commit time of snapshot that last updated this partition + private long lastUpdatedAt; + // ID of snapshot that last updated this partition + private long lastUpdatedSnapshotId; + + public enum Column { + PARTITION_DATA, + SPEC_ID, + DATA_RECORD_COUNT, + DATA_FILE_COUNT, + DATA_FILE_SIZE_IN_BYTES, + POSITION_DELETE_RECORD_COUNT, + POSITION_DELETE_FILE_COUNT, + EQUALITY_DELETE_RECORD_COUNT, + EQUALITY_DELETE_FILE_COUNT, + TOTAL_RECORD_COUNT, + LAST_UPDATED_AT, + LAST_UPDATED_SNAPSHOT_ID + } + + private PartitionEntry() {} + + public static Builder builder() { + return new Builder(); + } + + public PartitionData partitionData() { + return partitionData; + } + + public int specId() { + return specId; + } + + public long dataRecordCount() { + return dataRecordCount; + } + + public int dataFileCount() { + return dataFileCount; + } + + public long dataFileSizeInBytes() { + return dataFileSizeInBytes; + } + + public long posDeleteRecordCount() { + return posDeleteRecordCount; + } + + public int posDeleteFileCount() { + return posDeleteFileCount; + } + + public long eqDeleteRecordCount() { + return eqDeleteRecordCount; + } + + public int eqDeleteFileCount() { + return eqDeleteFileCount; + } + + public long totalRecordCount() { + return totalRecordCount; + } + + public long lastUpdatedAt() { + return lastUpdatedAt; + } + + public long lastUpdatedSnapshotId() { + return lastUpdatedSnapshotId; + } + + @Override + public void put(int i, Object v) { + switch (i) { + case 0: + this.partitionData = (PartitionData) v; + return; + case 1: + this.specId = (int) v; + return; + case 2: + this.dataRecordCount = (long) v; + return; + case 3: + this.dataFileCount = (int) v; + return; + case 4: + this.dataFileSizeInBytes = (long) v; + return; + case 5: + this.posDeleteRecordCount = (long) v; + return; + case 6: + this.posDeleteFileCount = (int) v; + return; + case 7: + this.eqDeleteRecordCount = (long) v; + return; + case 8: + this.eqDeleteFileCount = (int) v; + return; + case 9: + this.totalRecordCount = (long) v; + return; + case 10: + this.lastUpdatedAt = (long) v; + return; + case 11: + this.lastUpdatedSnapshotId = (long) v; + return; + default: + throw new UnsupportedOperationException("Unknown field ordinal: " + i); + } + } + + @Override + public Object get(int i) { + switch (i) { + case 0: + return partitionData; + case 1: + return specId; + case 2: + return dataRecordCount; + case 3: + return dataFileCount; + case 4: + return dataFileSizeInBytes; + case 5: + return posDeleteRecordCount; + case 6: + return posDeleteFileCount; + case 7: + return eqDeleteRecordCount; + case 8: + return eqDeleteFileCount; + case 9: + return totalRecordCount; + case 10: + return lastUpdatedAt; + case 11: + return lastUpdatedSnapshotId; + default: + throw new UnsupportedOperationException("Unknown field ordinal: " + i); + } + } + + @Override + public Schema getSchema() { + return prepareAvroSchema(partitionData.getPartitionType()); + } + + @Override + @SuppressWarnings("checkstyle:CyclomaticComplexity") + public boolean equals(Object o) { + if (this == o) { + return true; + } else if (!(o instanceof PartitionEntry)) { + return false; + } + + PartitionEntry that = (PartitionEntry) o; + return partitionData.equals(that.partitionData) + && specId == that.specId + && dataRecordCount == that.dataRecordCount + && dataFileCount == that.dataFileCount + && dataFileSizeInBytes == that.dataFileSizeInBytes + && posDeleteRecordCount == that.posDeleteRecordCount + && posDeleteFileCount == that.posDeleteFileCount + && eqDeleteRecordCount == that.eqDeleteRecordCount + && eqDeleteFileCount == that.eqDeleteFileCount + && totalRecordCount == that.totalRecordCount + && lastUpdatedAt == that.lastUpdatedAt + && lastUpdatedSnapshotId == that.lastUpdatedSnapshotId; + } + + @Override + public int hashCode() { + return Objects.hash( + partitionData, + specId, + dataRecordCount, + dataFileCount, + dataFileSizeInBytes, + posDeleteRecordCount, + posDeleteFileCount, + eqDeleteRecordCount, + eqDeleteFileCount, + totalRecordCount, + lastUpdatedAt, + lastUpdatedSnapshotId); + } + + public static org.apache.iceberg.Schema icebergSchema(Types.StructType partitionType) { + if (partitionType.fields().isEmpty()) { + throw new IllegalArgumentException("getting schema for an unpartitioned table"); + } + + return new org.apache.iceberg.Schema( + Types.NestedField.required(1, Column.PARTITION_DATA.name(), partitionType), + Types.NestedField.required(2, Column.SPEC_ID.name(), Types.IntegerType.get()), + Types.NestedField.required(3, Column.DATA_RECORD_COUNT.name(), Types.LongType.get()), + Types.NestedField.required(4, Column.DATA_FILE_COUNT.name(), Types.IntegerType.get()), + Types.NestedField.required(5, Column.DATA_FILE_SIZE_IN_BYTES.name(), Types.LongType.get()), + Types.NestedField.optional( + 6, Column.POSITION_DELETE_RECORD_COUNT.name(), Types.LongType.get()), + Types.NestedField.optional( + 7, Column.POSITION_DELETE_FILE_COUNT.name(), Types.IntegerType.get()), + Types.NestedField.optional( + 8, Column.EQUALITY_DELETE_RECORD_COUNT.name(), Types.LongType.get()), + Types.NestedField.optional( + 9, Column.EQUALITY_DELETE_FILE_COUNT.name(), Types.IntegerType.get()), + Types.NestedField.optional(10, Column.TOTAL_RECORD_COUNT.name(), Types.LongType.get()), + Types.NestedField.optional(11, Column.LAST_UPDATED_AT.name(), Types.LongType.get()), + Types.NestedField.optional( + 12, Column.LAST_UPDATED_SNAPSHOT_ID.name(), Types.LongType.get())); + } + + private static Schema prepareAvroSchema(Types.StructType partitionType) { + return AvroSchemaUtil.convert(icebergSchema(partitionType), "partitionEntry"); + } + + public static class Builder { + private PartitionData partitionData; + private int specId; + private long dataRecordCount; + private int dataFileCount; + private long dataFileSizeInBytes; + private long posDeleteRecordCount; + private int posDeleteFileCount; + private long eqDeleteRecordCount; + private int eqDeleteFileCount; + private long totalRecordCount; + private long lastUpdatedAt; + private long lastUpdatedSnapshotId; + + private Builder() {} + + public Builder withPartitionData(PartitionData newPartitionData) { + this.partitionData = newPartitionData; + return this; + } + + public Builder withSpecId(int newSpecId) { + this.specId = newSpecId; + return this; + } + + public Builder withDataRecordCount(long newDataRecordCount) { + this.dataRecordCount = newDataRecordCount; + return this; + } + + public Builder withDataFileCount(int newDataFileCount) { + this.dataFileCount = newDataFileCount; + return this; + } + + public Builder withDataFileSizeInBytes(long newDataFileSizeInBytes) { + this.dataFileSizeInBytes = newDataFileSizeInBytes; + return this; + } + + public Builder withPosDeleteRecordCount(Long newPosDeleteRecordCount) { + this.posDeleteRecordCount = newPosDeleteRecordCount; + return this; + } + + public Builder withPosDeleteFileCount(Integer newPosDeleteFileCount) { + this.posDeleteFileCount = newPosDeleteFileCount; + return this; + } + + public Builder withEqDeleteRecordCount(Long newEqDeleteRecordCount) { + this.eqDeleteRecordCount = newEqDeleteRecordCount; + return this; + } + + public Builder withEqDeleteFileCount(Integer newEqDeleteFileCount) { + this.eqDeleteFileCount = newEqDeleteFileCount; + return this; + } + + public Builder withTotalRecordCount(Long newTotalRecordCount) { + this.totalRecordCount = newTotalRecordCount; + return this; + } + + public Builder withLastUpdatedAt(Long newLastUpdatedAt) { + this.lastUpdatedAt = newLastUpdatedAt; + return this; + } + + public Builder withLastUpdatedSnapshotId(Long newLastUpdatedSnapshotId) { + this.lastUpdatedSnapshotId = newLastUpdatedSnapshotId; + return this; + } + + public PartitionEntry newInstance() { + return new PartitionEntry(); + } + + public PartitionEntry build() { + PartitionEntry partition = new PartitionEntry(); + partition.partitionData = partitionData; + partition.specId = specId; + partition.dataRecordCount = dataRecordCount; + partition.dataFileCount = dataFileCount; + partition.dataFileSizeInBytes = dataFileSizeInBytes; + partition.posDeleteRecordCount = posDeleteRecordCount; + partition.posDeleteFileCount = posDeleteFileCount; + partition.eqDeleteRecordCount = eqDeleteRecordCount; + partition.eqDeleteFileCount = eqDeleteFileCount; + partition.totalRecordCount = totalRecordCount; + partition.lastUpdatedAt = lastUpdatedAt; + partition.lastUpdatedSnapshotId = lastUpdatedSnapshotId; + return partition; + } + } +} diff --git a/core/src/main/java/org/apache/iceberg/Partitioning.java b/core/src/main/java/org/apache/iceberg/Partitioning.java index 7e4fcae333d8..872eb3bb09af 100644 --- a/core/src/main/java/org/apache/iceberg/Partitioning.java +++ b/core/src/main/java/org/apache/iceberg/Partitioning.java @@ -238,7 +238,10 @@ public static StructType groupingKeyType(Schema schema, Collection specs = table.specs().values(); + return partitionType(table.specs().values()); + } + + public static StructType partitionType(Collection specs) { return buildPartitionProjectionType("table partition", specs, allFieldIds(specs)); } diff --git a/core/src/test/java/org/apache/iceberg/TestPartitionEntry.java b/core/src/test/java/org/apache/iceberg/TestPartitionEntry.java new file mode 100644 index 000000000000..cb9fe9c9149c --- /dev/null +++ b/core/src/test/java/org/apache/iceberg/TestPartitionEntry.java @@ -0,0 +1,76 @@ +/* + * 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; + +import org.apache.iceberg.types.Types; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +public class TestPartitionEntry { + + @Test + public void testPartitionBuilder() { + Types.StructType partitionType = + Types.StructType.of( + Types.NestedField.required(1000, "field1", Types.StringType.get()), + Types.NestedField.required(1001, "field2", Types.IntegerType.get())); + PartitionData partitionData = new PartitionData(partitionType); + partitionData.set(0, "value1"); + partitionData.set(1, 42); + + PartitionEntry partition = + PartitionEntry.builder() + .withPartitionData(partitionData) + .withSpecId(123) + .withDataRecordCount(1000L) + .withDataFileCount(5) + .withDataFileSizeInBytes(1024L * 1024L) + .withPosDeleteRecordCount(50L) + .withPosDeleteFileCount(2) + .withEqDeleteRecordCount(20L) + .withEqDeleteFileCount(1) + .withTotalRecordCount(3000L) + .withLastUpdatedAt(1627900200L) + .withLastUpdatedSnapshotId(456789L) + .build(); + + // Verify the get method + Assertions.assertEquals(partitionData, partition.get(0)); + Assertions.assertEquals(123, partition.get(1)); + Assertions.assertEquals(1000L, partition.get(2)); + Assertions.assertEquals(5, partition.get(3)); + Assertions.assertEquals(1024L * 1024L, partition.get(4)); + Assertions.assertEquals(50L, partition.get(5)); + Assertions.assertEquals(2, partition.get(6)); + Assertions.assertEquals(20L, partition.get(7)); + Assertions.assertEquals(1, partition.get(8)); + Assertions.assertEquals(3000L, partition.get(9)); + Assertions.assertEquals(1627900200L, partition.get(10)); + Assertions.assertEquals(456789L, partition.get(11)); + + // Verify the put method + PartitionEntry newPartition = PartitionEntry.builder().newInstance(); + int size = partition.getSchema().getFields().size(); + for (int i = 0; i < size; i++) { + newPartition.put(i, partition.get(i)); + } + + Assertions.assertEquals(newPartition, partition); + } +} diff --git a/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java b/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java new file mode 100644 index 000000000000..83096088e442 --- /dev/null +++ b/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java @@ -0,0 +1,148 @@ +/* + * 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.data; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.Collection; +import java.util.Iterator; +import java.util.List; +import org.apache.avro.generic.GenericData; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionEntry; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Partitioning; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.TableOperations; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.DataWriter; +import org.apache.iceberg.io.InputFile; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.parquet.Parquet; +import org.apache.iceberg.parquet.ParquetAvroValueReaders; +import org.apache.iceberg.parquet.ParquetAvroWriter; +import org.apache.iceberg.types.Types; + +public final class PartitionStatsUtil { + + private PartitionStatsUtil() {} + + private static final String PARQUET_SUFFIX = ".parquet"; + + public static OutputFile newPartitionStatsFile( + TableOperations ops, long snapshotId, FileFormat format) { + return ops.io() + .newOutputFile( + ops.metadataFileLocation( + format.addExtension(String.format("partition-stats-%d", snapshotId)))); + } + + public static void writePartitionStatsFile( + Iterator partitions, OutputFile outputFile, Collection specs) { + validateFormat(outputFile.location()); + writeAsParquetFile( + PartitionEntry.icebergSchema(Partitioning.partitionType(specs)), partitions, outputFile); + } + + private static void validateFormat(String filePath) { + if (!filePath.toLowerCase().endsWith(PARQUET_SUFFIX)) { + throw new UnsupportedOperationException("Unsupported format : " + filePath); + } + } + + public static CloseableIterable readPartitionStatsFile( + Schema schema, InputFile inputFile) { + validateFormat(inputFile.location()); + // schema of partition column during read could be different from + // what is used for writing due to partition evolution. + // While reading, ParquetAvroValueReaders fills the data as per latest schema. + CloseableIterable records = + Parquet.read(inputFile) + .project(schema) + .createReaderFunc(fileSchema -> ParquetAvroValueReaders.buildReader(schema, fileSchema)) + .build(); + + return CloseableIterable.transform(records, record -> toPartition(schema, record)); + } + + private static PartitionEntry toPartition(Schema schema, GenericData.Record record) { + PartitionEntry partition = PartitionEntry.builder().newInstance(); + partition.put( + PartitionEntry.Column.PARTITION_DATA.ordinal(), + extractPartitionDataFromRecord(schema, record)); + + int recordCount = record.getSchema().getFields().size(); + for (int columnIndex = 1; columnIndex < recordCount; columnIndex++) { + partition.put(columnIndex, record.get(columnIndex)); + } + + return partition; + } + + private static PartitionData extractPartitionDataFromRecord( + Schema schema, GenericData.Record record) { + int partitionDataCount = + record + .getSchema() + .getField(PartitionEntry.Column.PARTITION_DATA.name()) + .schema() + .getFields() + .size(); + PartitionData partitionData = + new PartitionData( + (Types.StructType) + schema.findField(PartitionEntry.Column.PARTITION_DATA.name()).type()); + for (int partitionColIndex = 0; partitionColIndex < partitionDataCount; partitionColIndex++) { + partitionData.set( + partitionColIndex, + ((GenericData.Record) record.get(PartitionEntry.Column.PARTITION_DATA.ordinal())) + .get(partitionColIndex)); + } + + return partitionData; + } + + private static void writeAsParquetFile( + Schema schema, Iterator records, OutputFile outputFile) { + try (DataWriter dataWriter = + Parquet.writeData(outputFile) + .schema(schema) + .createWriterFunc(ParquetAvroWriter::buildWriter) + .overwrite() + .withSpec(PartitionSpec.unpartitioned()) + .withSortOrder(sortOrder(schema)) + .build()) { + records.forEachRemaining(dataWriter::write); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + } + + private static SortOrder sortOrder(Schema schema) { + SortOrder.Builder builder = SortOrder.builderFor(schema); + List partitionFields = + ((Types.StructType) schema.asStruct().fields().get(0).type()).fields(); + partitionFields.forEach( + field -> builder.asc(PartitionEntry.Column.PARTITION_DATA.name() + "." + field.name())); + + return builder.build(); + } +} diff --git a/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsUtil.java b/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsUtil.java new file mode 100644 index 000000000000..f7eafaaa3017 --- /dev/null +++ b/data/src/test/java/org/apache/iceberg/data/TestPartitionStatsUtil.java @@ -0,0 +1,125 @@ +/* + * 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.data; + +import java.nio.file.Paths; +import java.util.List; +import java.util.Random; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.Files; +import org.apache.iceberg.HasTableOperations; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionEntry; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Partitioning; +import org.apache.iceberg.Schema; +import org.apache.iceberg.SortOrder; +import org.apache.iceberg.Table; +import org.apache.iceberg.TestTables; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.types.Types; +import org.assertj.core.api.Assertions; +import org.junit.Rule; +import org.junit.Test; +import org.junit.rules.TemporaryFolder; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +public class TestPartitionStatsUtil { + private static final Logger LOG = LoggerFactory.getLogger(TestPartitionStatsUtil.class); + + private static final Schema SCHEMA = + new Schema( + Types.NestedField.required(1, "id", Types.LongType.get()), + Types.NestedField.optional(2, "data", Types.StringType.get()), + Types.NestedField.optional(3, "binary", Types.BinaryType.get())); + + @Rule public TemporaryFolder temp = new TemporaryFolder(); + + @Test + public void testPartitionStats() throws Exception { + PartitionSpec spec = PartitionSpec.builderFor(SCHEMA).identity("id").identity("binary").build(); + Table testTable = + TestTables.create( + temp.newFolder("test_partition_stats"), + "test_partition_stats", + SCHEMA, + spec, + SortOrder.unsorted(), + 2); + + Schema schema = + PartitionEntry.icebergSchema(Partitioning.partitionType(testTable.specs().values())); + + ImmutableList.Builder partitionListBuilder = ImmutableList.builder(); + + long seed = System.currentTimeMillis(); + LOG.info("Seed used for random generator is {}", seed); + Random random = new Random(seed); + + for (int i = 0; i < 42; i++) { + PartitionData partitionData = + new PartitionData( + schema.findField(PartitionEntry.Column.PARTITION_DATA.name()).type().asStructType()); + partitionData.set(0, random.nextLong()); + + PartitionEntry partition = + PartitionEntry.builder() + .withPartitionData(partitionData) + .withSpecId(random.nextInt(10)) + .withDataRecordCount(random.nextLong()) + .withDataFileCount(random.nextInt()) + .withDataFileSizeInBytes(1024L * random.nextInt(20)) + .withPosDeleteRecordCount(random.nextLong()) + .withPosDeleteFileCount(random.nextInt()) + .withEqDeleteRecordCount(random.nextLong()) + .withEqDeleteFileCount(random.nextInt()) + .withTotalRecordCount(random.nextLong()) + .withLastUpdatedAt(random.nextLong()) + .withLastUpdatedSnapshotId(random.nextLong()) + .build(); + + partitionListBuilder.add(partition); + } + List records = partitionListBuilder.build(); + + OutputFile outputFile = + PartitionStatsUtil.newPartitionStatsFile( + ((HasTableOperations) testTable).operations(), 42L, FileFormat.PARQUET); + PartitionStatsUtil.writePartitionStatsFile( + records.iterator(), outputFile, testTable.specs().values()); + + Assertions.assertThat(Paths.get(outputFile.location())).exists(); + + List rows; + try (CloseableIterable recordIterator = + PartitionStatsUtil.readPartitionStatsFile( + schema, Files.localInput(outputFile.location()))) { + rows = Lists.newArrayList(recordIterator); + } + + Assertions.assertThat(rows).hasSize(records.size()); + for (int i = 0; i < records.size(); i++) { + Assertions.assertThat(rows.get(i)).isEqualTo(records.get(i)); + } + } +} From 36fc685908b4ac8cec11a53f1ca4112a7e07e160 Mon Sep 17 00:00:00 2001 From: ajantha-bhat Date: Thu, 21 Dec 2023 21:41:52 +0530 Subject: [PATCH 2/2] Spark 3.5: Spark action to compute the partition stats --- .../iceberg/actions/ActionsProvider.java | 6 + .../actions/ComputePartitionStats.java | 35 +++ .../org/apache/iceberg/PartitionEntry.java | 99 +++++++++ .../actions/BaseComputePartitionStats.java | 33 +++ .../iceberg/data/PartitionStatsUtil.java | 5 +- .../spark/actions/BaseSparkAction.java | 68 ++++++ .../ComputePartitionStatsSparkAction.java | 202 ++++++++++++++++++ .../spark/actions/PartitionEntryBean.java | 116 ++++++++++ .../iceberg/spark/actions/SparkActions.java | 5 + .../TestComputePartitionStatsAction.java | 189 ++++++++++++++++ .../TestComputePartitionStatsActionPerf.java | 153 +++++++++++++ 11 files changed, 910 insertions(+), 1 deletion(-) create mode 100644 api/src/main/java/org/apache/iceberg/actions/ComputePartitionStats.java create mode 100644 core/src/main/java/org/apache/iceberg/actions/BaseComputePartitionStats.java create mode 100644 spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/ComputePartitionStatsSparkAction.java create mode 100644 spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/PartitionEntryBean.java create mode 100644 spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsAction.java create mode 100644 spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsActionPerf.java diff --git a/api/src/main/java/org/apache/iceberg/actions/ActionsProvider.java b/api/src/main/java/org/apache/iceberg/actions/ActionsProvider.java index 2d6ff2679a17..d40a38867b90 100644 --- a/api/src/main/java/org/apache/iceberg/actions/ActionsProvider.java +++ b/api/src/main/java/org/apache/iceberg/actions/ActionsProvider.java @@ -70,4 +70,10 @@ default RewritePositionDeleteFiles rewritePositionDeletes(Table table) { throw new UnsupportedOperationException( this.getClass().getName() + " does not implement rewritePositionDeletes"); } + + /** Instantiates an action to compute partition statistics and register it to table metadata. */ + default ComputePartitionStats computePartitionStatistics(Table table) { + throw new UnsupportedOperationException( + this.getClass().getName() + " does not implement computePartitionStatistics"); + } } diff --git a/api/src/main/java/org/apache/iceberg/actions/ComputePartitionStats.java b/api/src/main/java/org/apache/iceberg/actions/ComputePartitionStats.java new file mode 100644 index 000000000000..bb6c88c56d4a --- /dev/null +++ b/api/src/main/java/org/apache/iceberg/actions/ComputePartitionStats.java @@ -0,0 +1,35 @@ +/* + * 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.actions; + +import org.apache.iceberg.PartitionStatisticsFile; + +/** An action to compute and register partition stats. */ +public interface ComputePartitionStats + extends Action { + + /** The action result that contains a summary of the execution. */ + interface Result { + /** + * Returns the output file which is registered to the table metadata, null if the table is + * non-partitioned or empty. + */ + PartitionStatisticsFile outputFile(); + } +} diff --git a/core/src/main/java/org/apache/iceberg/PartitionEntry.java b/core/src/main/java/org/apache/iceberg/PartitionEntry.java index ac6da58c5d5a..677c9d484a62 100644 --- a/core/src/main/java/org/apache/iceberg/PartitionEntry.java +++ b/core/src/main/java/org/apache/iceberg/PartitionEntry.java @@ -18,11 +18,15 @@ */ package org.apache.iceberg; +import java.util.List; +import java.util.Map; import java.util.Objects; import org.apache.avro.Schema; import org.apache.avro.generic.IndexedRecord; import org.apache.iceberg.avro.AvroSchemaUtil; +import org.apache.iceberg.io.CloseableIterable; import org.apache.iceberg.types.Types; +import org.apache.iceberg.util.PartitionUtil; public class PartitionEntry implements IndexedRecord { private PartitionData partitionData; @@ -261,6 +265,24 @@ private static Schema prepareAvroSchema(Types.StructType partitionType) { return AvroSchemaUtil.convert(icebergSchema(partitionType), "partitionEntry"); } + public synchronized PartitionEntry update(PartitionEntry entry) { + this.specId = Math.max(this.specId, entry.specId); + this.dataRecordCount += entry.dataRecordCount; + this.dataFileCount += entry.dataFileCount; + this.dataFileSizeInBytes += entry.dataFileSizeInBytes; + this.posDeleteRecordCount += entry.posDeleteRecordCount; + this.posDeleteFileCount += entry.posDeleteFileCount; + this.eqDeleteRecordCount += entry.eqDeleteRecordCount; + this.eqDeleteFileCount += entry.eqDeleteFileCount; + this.totalRecordCount += entry.totalRecordCount; + if (this.lastUpdatedAt < entry.lastUpdatedAt) { + this.lastUpdatedAt = entry.lastUpdatedAt(); + this.lastUpdatedSnapshotId = entry.lastUpdatedSnapshotId; + } + + return this; + } + public static class Builder { private PartitionData partitionData; private int specId; @@ -358,4 +380,81 @@ public PartitionEntry build() { return partition; } } + + public static CloseableIterable fromManifest(Table table, ManifestFile manifest) { + CloseableIterable>> entries = + CloseableIterable.transform( + ManifestFiles.open(manifest, table.io(), table.specs()) + .select(scanColumns(manifest.content())) // don't select stats columns + .liveEntries(), + t -> + (ManifestEntry>) + // defensive copy of manifest entry without stats columns + t.copyWithoutStats()); + + Types.StructType partitionType = Partitioning.partitionType(table); + return CloseableIterable.transform( + entries, entry -> fromManifestEntry(entry, table, partitionType)); + } + + private static PartitionEntry fromManifestEntry( + ManifestEntry entry, Table table, Types.StructType partitionType) { + PartitionEntry.Builder builder = PartitionEntry.builder(); + builder + .withSpecId(entry.file().specId()) + .withPartitionData(coercedPartitionData(entry.file(), table.specs(), partitionType)); + Snapshot snapshot = table.snapshot(entry.snapshotId()); + if (snapshot != null) { + builder + .withLastUpdatedSnapshotId(snapshot.snapshotId()) + .withLastUpdatedAt(snapshot.timestampMillis()); + } + + switch (entry.file().content()) { + case DATA: + builder + .withDataFileCount(1) + .withDataRecordCount(entry.file().recordCount()) + .withDataFileSizeInBytes(entry.file().fileSizeInBytes()); + break; + case POSITION_DELETES: + builder.withPosDeleteFileCount(1).withPosDeleteRecordCount(entry.file().recordCount()); + break; + case EQUALITY_DELETES: + builder.withEqDeleteFileCount(1).withEqDeleteRecordCount(entry.file().recordCount()); + break; + default: + throw new UnsupportedOperationException( + "Unsupported file content type: " + entry.file().content()); + } + + // TODO: optionally compute TOTAL_RECORD_COUNT based on the flag + return builder.build(); + } + + private static PartitionData coercedPartitionData( + ContentFile file, Map specs, Types.StructType partitionType) { + // keep the partition data as per the unified spec by coercing + StructLike partition = + PartitionUtil.coercePartition(partitionType, specs.get(file.specId()), file.partition()); + PartitionData data = new PartitionData(partitionType); + for (int i = 0; i < partitionType.fields().size(); i++) { + Object val = partition.get(i, partitionType.fields().get(i).type().typeId().javaClass()); + if (val != null) { + data.set(i, val); + } + } + return data; + } + + private static List scanColumns(ManifestContent content) { + switch (content) { + case DATA: + return BaseScan.SCAN_COLUMNS; + case DELETES: + return BaseScan.DELETE_SCAN_COLUMNS; + default: + throw new UnsupportedOperationException("Cannot read unknown manifest type: " + content); + } + } } diff --git a/core/src/main/java/org/apache/iceberg/actions/BaseComputePartitionStats.java b/core/src/main/java/org/apache/iceberg/actions/BaseComputePartitionStats.java new file mode 100644 index 000000000000..80f5876527b9 --- /dev/null +++ b/core/src/main/java/org/apache/iceberg/actions/BaseComputePartitionStats.java @@ -0,0 +1,33 @@ +/* + * 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.actions; + +import org.immutables.value.Value; + +@Value.Enclosing +@SuppressWarnings("ImmutablesStyle") +@Value.Style( + typeImmutableEnclosing = "ImmutableComputePartitionStats", + visibilityString = "PUBLIC", + builderVisibilityString = "PUBLIC") +interface BaseComputePartitionStats extends ComputePartitionStats { + + @Value.Immutable + interface Result extends ComputePartitionStats.Result {} +} diff --git a/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java b/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java index 83096088e442..4441eef5cb08 100644 --- a/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java +++ b/data/src/main/java/org/apache/iceberg/data/PartitionStatsUtil.java @@ -23,6 +23,7 @@ import java.util.Collection; import java.util.Iterator; import java.util.List; +import java.util.UUID; import org.apache.avro.generic.GenericData; import org.apache.iceberg.FileFormat; import org.apache.iceberg.PartitionData; @@ -49,10 +50,12 @@ private PartitionStatsUtil() {} public static OutputFile newPartitionStatsFile( TableOperations ops, long snapshotId, FileFormat format) { + // TODO: UUID is temp, remove it. return ops.io() .newOutputFile( ops.metadataFileLocation( - format.addExtension(String.format("partition-stats-%d", snapshotId)))); + format.addExtension( + String.format("partition-stats-%s-%d", UUID.randomUUID(), snapshotId)))); } public static void writePartitionStatsFile( diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/BaseSparkAction.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/BaseSparkAction.java index 53ce7418f3ec..2b6d3988a049 100644 --- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/BaseSparkAction.java +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/BaseSparkAction.java @@ -19,8 +19,13 @@ package org.apache.iceberg.spark.actions; import static org.apache.iceberg.MetadataTableType.ALL_MANIFESTS; +import static org.apache.iceberg.MetadataTableType.ENTRIES; import static org.apache.spark.sql.functions.col; +import static org.apache.spark.sql.functions.first; import static org.apache.spark.sql.functions.lit; +import static org.apache.spark.sql.functions.max; +import static org.apache.spark.sql.functions.sum; +import static org.apache.spark.sql.functions.when; import java.util.Collection; import java.util.Iterator; @@ -42,6 +47,7 @@ import org.apache.iceberg.MetadataTableType; import org.apache.iceberg.PartitionSpec; import org.apache.iceberg.ReachableFileUtil; +import org.apache.iceberg.Snapshot; import org.apache.iceberg.StaticTableOperations; import org.apache.iceberg.Table; import org.apache.iceberg.TableMetadata; @@ -72,6 +78,8 @@ import org.apache.spark.sql.Dataset; import org.apache.spark.sql.Row; import org.apache.spark.sql.SparkSession; +import org.apache.spark.sql.functions; +import org.apache.spark.sql.types.DataTypes; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -164,6 +172,66 @@ protected Dataset contentFileDS(Table table, Set snapshotIds) { return manifestBeanDS.flatMap(new ReadManifest(tableBroadcast), FileInfo.ENCODER); } + protected Dataset partitionEntryDS(Table table) { + Dataset dataset = + loadMetadataTable(table, ENTRIES) + .filter(col("status").$less(2)) + .select( + col("data_file.spec_id").as("SPEC_ID"), + col("data_file.partition").as("PARTITION_DATA"), + when(col("data_file.content").equalTo(0), col("data_file.record_count")) + .otherwise(lit(0)) + .as("DATA_RECORD_COUNT"), + when(col("data_file.content").equalTo(0), lit(1)) + .otherwise(lit(0)) + .as("DATA_FILE_COUNT"), + when(col("data_file.content").equalTo(0), col("data_file.file_size_in_bytes")) + .otherwise(lit(0)) + .as("DATA_FILE_SIZE_IN_BYTES"), + when(col("data_file.content").equalTo(1), col("data_file.record_count")) + .otherwise(lit(0)) + .as("POSITION_DELETE_RECORD_COUNT"), + when(col("data_file.content").equalTo(1), lit(1)) + .otherwise(lit(0)) + .as("POSITION_DELETE_FILE_COUNT"), + when(col("data_file.content").equalTo(2), col("data_file.record_count")) + .otherwise(lit(0)) + .as("EQUALITY_DELETE_RECORD_COUNT"), + when(col("data_file.content").equalTo(2), lit(1)) + .otherwise(lit(0)) + .as("EQUALITY_DELETE_FILE_COUNT"), + functions + .udf( + (Long snapshotId) -> lastUpdatedTime(snapshotId, table), DataTypes.LongType) + .apply(col("snapshot_id")) + .as("LAST_UPDATED_AT"), + col("snapshot_id").as("LAST_UPDATED_SNAPSHOT_ID"), + lit(0) + .alias("TOTAL_RECORD_COUNT")); // TODO: not sure if this can be computed by this + // distributed algorithm. This was meant to be + // effective count after applying deletes. + + return dataset + .groupBy(col("PARTITION_DATA")) + .agg( + max(col("LAST_UPDATED_SNAPSHOT_ID")).as("LAST_UPDATED_SNAPSHOT_ID"), + first(col("LAST_UPDATED_AT")).as("LAST_UPDATED_AT"), + max(col("SPEC_ID")).as("SPEC_ID"), + sum(col("DATA_FILE_COUNT")).as("DATA_FILE_COUNT"), + sum(col("DATA_RECORD_COUNT")).as("DATA_RECORD_COUNT"), + sum(col("DATA_FILE_SIZE_IN_BYTES")).as("DATA_FILE_SIZE_IN_BYTES"), + sum(col("POSITION_DELETE_FILE_COUNT")).as("POSITION_DELETE_FILE_COUNT"), + sum(col("POSITION_DELETE_RECORD_COUNT")).as("POSITION_DELETE_RECORD_COUNT"), + sum(col("EQUALITY_DELETE_FILE_COUNT")).as("EQUALITY_DELETE_FILE_COUNT"), + sum(col("EQUALITY_DELETE_RECORD_COUNT")).as("EQUALITY_DELETE_RECORD_COUNT"), + sum(col("TOTAL_RECORD_COUNT")).as("TOTAL_RECORD_COUNT")); + } + + public static long lastUpdatedTime(long snapshotId, Table table) { + Snapshot snapshot = table.snapshot(snapshotId); + return snapshot == null ? 0 : snapshot.timestampMillis(); + } + protected Dataset manifestDS(Table table) { return manifestDS(table, null); } diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/ComputePartitionStatsSparkAction.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/ComputePartitionStatsSparkAction.java new file mode 100644 index 000000000000..ef545a441236 --- /dev/null +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/ComputePartitionStatsSparkAction.java @@ -0,0 +1,202 @@ +/* + * 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.actions; + +import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT; +import static org.apache.iceberg.TableProperties.DEFAULT_FILE_FORMAT_DEFAULT; + +import java.io.File; +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import org.apache.commons.io.FileUtils; +import org.apache.iceberg.BaseTable; +import org.apache.iceberg.FileFormat; +import org.apache.iceberg.ImmutableGenericPartitionStatisticsFile; +import org.apache.iceberg.ManifestFile; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionEntry; +import org.apache.iceberg.PartitionStatisticsFile; +import org.apache.iceberg.Table; +import org.apache.iceberg.actions.ComputePartitionStats; +import org.apache.iceberg.actions.ImmutableComputePartitionStats; +import org.apache.iceberg.data.PartitionStatsUtil; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.io.OutputFile; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.spark.JobGroupInfo; +import org.apache.iceberg.util.Tasks; +import org.apache.iceberg.util.ThreadPools; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.SparkSession; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * An implementation of {@link ComputePartitionStats} that computes and registers the partition + * stats to table metadata + */ +public class ComputePartitionStatsSparkAction + extends BaseSparkAction implements ComputePartitionStats { + private static final Logger LOG = LoggerFactory.getLogger(ComputePartitionStatsSparkAction.class); + private final Table table; + + private boolean local = false; + + ComputePartitionStatsSparkAction(SparkSession spark, Table table) { + super(spark); + this.table = table; + } + + public ComputePartitionStatsSparkAction localCompute(boolean localCompute) { + this.local = localCompute; + return this; + } + + @Override + protected ComputePartitionStatsSparkAction self() { + return this; + } + + @Override + public Result execute() { + String jobDesc = String.format("Computing partition stats for the table %s", table.name()); + JobGroupInfo info = newJobGroupInfo("COMPUTE-PARTITION-STATS", jobDesc); + return withJobGroupInfo(info, this::doExecute); + } + + private Result doExecute() { + long currentSnapshotId = table.currentSnapshot().snapshotId(); + if (currentSnapshotId == -1) { + // when the action is executed on an empty table. + return null; + } + + FileFormat fileFormat = + FileFormat.fromString( + table.properties().getOrDefault(DEFAULT_FILE_FORMAT, DEFAULT_FILE_FORMAT_DEFAULT)); + OutputFile outputFile = + PartitionStatsUtil.newPartitionStatsFile( + ((BaseTable) table).operations(), currentSnapshotId, fileFormat); + if (local) { + localCompute(outputFile); + } else { + // don't have control over output file name from Spark. hence writing to a temp location + // and moving stats file to the metadata folder after that. + String tempLocation = table.location() + "/metadata/temp" + UUID.randomUUID(); + distributedCompute(tempLocation); + File[] files = + new File(tempLocation.replaceFirst("file:", "")) + .listFiles((dir, name) -> name.toLowerCase().endsWith(".parquet")); + if (files == null || files.length == 0) { + LOG.error("partition stats file not found in temp location {}", tempLocation); + return null; + } + + try { + // Since coalesce(1) is used, only one file will be written. + files[0].renameTo(new File(outputFile.location().replaceFirst("file:", ""))); + FileUtils.deleteDirectory(new File(tempLocation.replaceFirst("file:", ""))); + } catch (IOException e) { + throw new RuntimeException(e); + } + } + + PartitionStatisticsFile statisticsFile = + ImmutableGenericPartitionStatisticsFile.builder() + .snapshotId(currentSnapshotId) + .path(outputFile.location()) + .fileSizeInBytes(outputFile.toInputFile().getLength()) + .build(); + table.updatePartitionStatistics().setPartitionStatistics(statisticsFile).commit(); + LOG.info( + "Registered the partition stats file: {} for the snapshot id {} to the table {}", + outputFile.location(), + currentSnapshotId, + table.name()); + + return ImmutableComputePartitionStats.Result.builder().outputFile(statisticsFile).build(); + } + + private void distributedCompute(String outputDir) { + Dataset dataset = partitionEntryDS(table); + dataset + .select( + "PARTITION_DATA", + "SPEC_ID", + "DATA_RECORD_COUNT", + "DATA_FILE_COUNT", + "DATA_FILE_SIZE_IN_BYTES", + "POSITION_DELETE_RECORD_COUNT", + "POSITION_DELETE_FILE_COUNT", + "EQUALITY_DELETE_RECORD_COUNT", + "EQUALITY_DELETE_FILE_COUNT", + "TOTAL_RECORD_COUNT", + "LAST_UPDATED_AT", + "LAST_UPDATED_SNAPSHOT_ID") + .coalesce(1) + .write() + .format("parquet") + .mode("overwrite") + .option("path", outputDir) + .save(); + } + + private void localCompute(OutputFile outputFile) { + Map partitionEntryMap = Maps.newConcurrentMap(); + List manifestFiles = table.currentSnapshot().allManifests(table.io()); + + Tasks.foreach(manifestFiles) + .retry(3) + .suppressFailureWhenFinished() + .executeWith(ThreadPools.getWorkerPool()) + .onFailure( + (file, thrown) -> + LOG.warn( + "Failed to compute the partition stats for the manifest file: {}", + file.path(), + thrown)) + .run( + manifest -> { + try (CloseableIterable entries = + PartitionEntry.fromManifest(table, manifest)) { + entries.forEach( + entry -> + partitionEntryMap.compute( + entry.partitionData(), + (key, existingEntry) -> { + if (existingEntry != null) { + existingEntry.update(entry); + return existingEntry; + } else { + return entry; + } + })); + } catch (IOException e) { + throw new UncheckedIOException(e); + } + }); + + PartitionStatsUtil.writePartitionStatsFile( + partitionEntryMap.values().iterator(), outputFile, table.specs().values()); + } +} diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/PartitionEntryBean.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/PartitionEntryBean.java new file mode 100644 index 000000000000..49a0836c8105 --- /dev/null +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/PartitionEntryBean.java @@ -0,0 +1,116 @@ +/* + * 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.actions; + +import org.apache.iceberg.FileContent; +import org.apache.spark.sql.Encoder; +import org.apache.spark.sql.Encoders; + +public class PartitionEntryBean { + public static final Encoder ENCODER = Encoders.bean(PartitionEntryBean.class); + + private FileContent content; + private int specId; + // Using Json string representation instead of `PartitionData` object due to below error. + // Caused by: org.apache.spark.SparkUnsupportedOperationException: Cannot have circular + // references in bean class, but got the circular reference of class class + // org.apache.avro.Schema. + private String partition; + private long recordCount; + private long fileSizeInBytes; + private long lastUpdatedAt; + private long lastUpdatedSnapshotId; + + public PartitionEntryBean( + FileContent content, + int specId, + String partition, + long recordCount, + long fileSizeInBytes, + long lastUpdatedAt, + long lastUpdatedSnapshotId) { + this.content = content; + this.specId = specId; + this.partition = partition; + this.recordCount = recordCount; + this.fileSizeInBytes = fileSizeInBytes; + this.lastUpdatedAt = lastUpdatedAt; + this.lastUpdatedSnapshotId = lastUpdatedSnapshotId; + } + + public PartitionEntryBean() {} + + // Note: Getter and Setter naming conventions has to be as per Java standard for Spark codegen. + + public FileContent getContent() { + return content; + } + + public void setContent(FileContent content) { + this.content = content; + } + + public int getSpecId() { + return specId; + } + + public void setSpecId(int specId) { + this.specId = specId; + } + + public String getPartition() { + return partition; + } + + public void setPartition(String partition) { + this.partition = partition; + } + + public long getRecordCount() { + return recordCount; + } + + public void setRecordCount(long recordCount) { + this.recordCount = recordCount; + } + + public long getFileSizeInBytes() { + return fileSizeInBytes; + } + + public void setFileSizeInBytes(long fileSizeInBytes) { + this.fileSizeInBytes = fileSizeInBytes; + } + + public long getLastUpdatedAt() { + return lastUpdatedAt; + } + + public void setLastUpdatedAt(long lastUpdatedAt) { + this.lastUpdatedAt = lastUpdatedAt; + } + + public long getLastUpdatedSnapshotId() { + return lastUpdatedSnapshotId; + } + + public void setLastUpdatedSnapshotId(long lastUpdatedSnapshotId) { + this.lastUpdatedSnapshotId = lastUpdatedSnapshotId; + } +} diff --git a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/SparkActions.java b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/SparkActions.java index fb67ded96e35..0c684d516005 100644 --- a/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/SparkActions.java +++ b/spark/v3.5/spark/src/main/java/org/apache/iceberg/spark/actions/SparkActions.java @@ -96,4 +96,9 @@ public DeleteReachableFilesSparkAction deleteReachableFiles(String metadataLocat public RewritePositionDeleteFilesSparkAction rewritePositionDeletes(Table table) { return new RewritePositionDeleteFilesSparkAction(spark, table); } + + @Override + public ComputePartitionStatsSparkAction computePartitionStatistics(Table table) { + return new ComputePartitionStatsSparkAction(spark, table); + } } diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsAction.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsAction.java new file mode 100644 index 000000000000..54bda0b34f49 --- /dev/null +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsAction.java @@ -0,0 +1,189 @@ +/* + * 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.actions; + +import static org.apache.iceberg.types.Types.NestedField.optional; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.util.Arrays; +import java.util.List; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.Files; +import org.apache.iceberg.Parameter; +import org.apache.iceberg.ParameterizedTestExtension; +import org.apache.iceberg.Parameters; +import org.apache.iceberg.PartitionData; +import org.apache.iceberg.PartitionEntry; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Partitioning; +import org.apache.iceberg.Schema; +import org.apache.iceberg.Table; +import org.apache.iceberg.actions.ComputePartitionStats; +import org.apache.iceberg.data.PartitionStatsUtil; +import org.apache.iceberg.hadoop.HadoopTables; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.spark.TestBase; +import org.apache.iceberg.spark.source.ThreeColumnRecord; +import org.apache.iceberg.types.Types; +import org.apache.spark.sql.Dataset; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.apache.spark.sql.catalyst.expressions.GenericRowWithSchema; +import org.assertj.core.api.Assertions; +import org.assertj.core.groups.Tuple; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.TestTemplate; +import org.junit.jupiter.api.extension.ExtendWith; +import org.junit.jupiter.api.io.TempDir; + +@ExtendWith(ParameterizedTestExtension.class) +public class TestComputePartitionStatsAction extends TestBase { + + private static final HadoopTables TABLES = new HadoopTables(new Configuration()); + private static final Schema SCHEMA = + new Schema( + optional(1, "c1", Types.IntegerType.get()), + optional(2, "c2", Types.StringType.get()), + optional(3, "c3", Types.StringType.get())); + + protected static final PartitionSpec SPEC = + PartitionSpec.builderFor(SCHEMA).identity("c2").identity("c3").build(); + + @Parameters(name = "localCompute = {0}") + protected static List parameters() { + return Arrays.asList(true, false); + } + + @Parameter private Boolean localCompute; + + private String tableLocation = null; + + @TempDir private Path temp; + + @BeforeEach + public void setupTableLocation() throws Exception { + File tableDir = temp.resolve("junit").toFile(); + this.tableLocation = tableDir.toURI().toString(); + } + + @TestTemplate + public void testPartitionTable() { + Table table = TABLES.create(SCHEMA, SPEC, Maps.newHashMap(), tableLocation); + + // foo, A -> 3 records + // foo, B -> 1 record + // bar, A -> 1 record + // bar, B -> 2 records + List records = + ImmutableList.of( + new ThreeColumnRecord(1, "foo", "A"), + new ThreeColumnRecord(2, "foo", "B"), + new ThreeColumnRecord(3, "foo", "A"), + new ThreeColumnRecord(4, "bar", "B"), + new ThreeColumnRecord(5, "bar", "A"), + new ThreeColumnRecord(6, "bar", "B"), + new ThreeColumnRecord(7, "foo", "A")); + + Dataset df = spark.createDataFrame(records, ThreeColumnRecord.class).coalesce(1); + + // insert twice + df.select("c1", "c2", "c3").write().format("iceberg").mode("append").save(tableLocation); + df.select("c1", "c2", "c3").write().format("iceberg").mode("append").save(tableLocation); + + List validFiles = + spark + .read() + .format("iceberg") + .load(tableLocation + "#files") + .select("file_path") + .as(Encoders.STRING()) + .collectAsList(); + Assertions.assertThat(validFiles).as("Should be 8 valid files. 4 files per insert").hasSize(8); + Assertions.assertThat(table.partitionStatisticsFiles()).isEmpty(); + + SparkActions actions = SparkActions.get(); + ComputePartitionStats.Result result = + actions.computePartitionStatistics(table).localCompute(localCompute).execute(); + Assertions.assertThat(table.partitionStatisticsFiles()).containsExactly(result.outputFile()); + + if (localCompute) { + // read the partition entries from the stats file + Schema schema = + PartitionEntry.icebergSchema(Partitioning.partitionType(table.specs().values())); + List rows; + try (CloseableIterable recordIterator = + PartitionStatsUtil.readPartitionStatsFile( + schema, Files.localInput(result.outputFile().path()))) { + rows = Lists.newArrayList(recordIterator); + } catch (IOException e) { + throw new RuntimeException(e); + } + + Types.StructType partitionType = + schema.findField(PartitionEntry.Column.PARTITION_DATA.name()).type().asStructType(); + Assertions.assertThat(rows) + .extracting( + PartitionEntry::partitionData, + PartitionEntry::dataRecordCount, + PartitionEntry::dataFileCount) + .containsExactlyInAnyOrder( + Tuple.tuple(partitionData(partitionType, "foo", "A"), 6L, 2), + Tuple.tuple(partitionData(partitionType, "foo", "B"), 2L, 2), + Tuple.tuple(partitionData(partitionType, "bar", "A"), 2L, 2), + Tuple.tuple(partitionData(partitionType, "bar", "B"), 4L, 2)); + } else { + // can't use PartitionStatsUtil.readPartitionStatsFile as it uses + // ParquetAvroValueReaders$ReadBuilder + // and since native parquet doesn't write column ids, reader throws NPE. + List rows = + spark + .read() + .parquet(result.outputFile().path()) + .select("PARTITION_DATA", "DATA_RECORD_COUNT", "DATA_FILE_COUNT") + .collectAsList(); + Assertions.assertThat(rows) + .extracting( + row -> ((GenericRowWithSchema) row.get(0)).values()[0], + row -> ((GenericRowWithSchema) row.get(0)).values()[1], + row -> row.getLong(1), + row -> row.getLong(2)) + .containsExactlyInAnyOrder( + Tuple.tuple("foo", "A", 6L, 2L), + Tuple.tuple("foo", "B", 2L, 2L), + Tuple.tuple("bar", "A", 2L, 2L), + Tuple.tuple("bar", "B", 4L, 2L)); + // TODO: while aggregating, few Int datatype auto casted to Long. Finalize the data types. + } + } + + private static PartitionData partitionData(Types.StructType partitionType, String c2, String c3) { + PartitionData partitionData = new PartitionData(partitionType); + partitionData.set(0, c2); + partitionData.set(1, c3); + return partitionData; + } + + // TODO: add the testcase with delete files (pos and equality) + // and also validate each and every field of partition stats. +} diff --git a/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsActionPerf.java b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsActionPerf.java new file mode 100644 index 000000000000..2936a076a522 --- /dev/null +++ b/spark/v3.5/spark/src/test/java/org/apache/iceberg/spark/actions/TestComputePartitionStatsActionPerf.java @@ -0,0 +1,153 @@ +/* + * 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.actions; + +import static org.apache.iceberg.types.Types.NestedField.optional; + +import java.io.File; +import java.io.IOException; +import java.nio.file.Path; +import java.util.List; +import org.apache.hadoop.conf.Configuration; +import org.apache.iceberg.AppendFiles; +import org.apache.iceberg.DataFile; +import org.apache.iceberg.FileGenerationUtil; +import org.apache.iceberg.Files; +import org.apache.iceberg.PartitionEntry; +import org.apache.iceberg.PartitionSpec; +import org.apache.iceberg.Partitioning; +import org.apache.iceberg.Schema; +import org.apache.iceberg.StructLike; +import org.apache.iceberg.Table; +import org.apache.iceberg.TestHelpers; +import org.apache.iceberg.actions.ComputePartitionStats; +import org.apache.iceberg.data.PartitionStatsUtil; +import org.apache.iceberg.hadoop.HadoopTables; +import org.apache.iceberg.io.CloseableIterable; +import org.apache.iceberg.relocated.com.google.common.collect.Lists; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.apache.iceberg.spark.TestBase; +import org.apache.iceberg.types.Types; +import org.apache.spark.sql.Encoders; +import org.apache.spark.sql.Row; +import org.assertj.core.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; + +public class TestComputePartitionStatsActionPerf extends TestBase { + + private static final HadoopTables TABLES = new HadoopTables(new Configuration()); + private static final Schema SCHEMA = + new Schema( + optional(1, "c1", Types.IntegerType.get()), + optional(2, "c2", Types.StringType.get()), + optional(3, "c3", Types.StringType.get())); + + protected static final PartitionSpec SPEC = + PartitionSpec.builderFor(SCHEMA).identity("c1").build(); + + private String tableLocation = null; + + @TempDir private Path temp; + + @BeforeEach + public void setupTableLocation() throws Exception { + File tableDir = temp.resolve("junit").toFile(); + this.tableLocation = tableDir.toURI().toString(); + } + + @Test + public void testPerf() { + Table table = TABLES.create(SCHEMA, SPEC, Maps.newHashMap(), tableLocation); + final int partitionCount = 20; + final int datafilesPerPartitionCount = 10000; + + for (int partitionOrdinal = 0; partitionOrdinal < partitionCount; partitionOrdinal++) { + StructLike partition = TestHelpers.Row.of(partitionOrdinal); + + AppendFiles appendFiles = table.newAppend(); + + for (int fileOrdinal = 0; fileOrdinal < datafilesPerPartitionCount; fileOrdinal++) { + DataFile dataFile = FileGenerationUtil.generateDataFile(table, partition); + appendFiles.appendFile(dataFile); + } + + appendFiles.commit(); + } + + List validFiles = + spark + .read() + .format("iceberg") + .load(tableLocation + "#files") + .select("file_path") + .as(Encoders.STRING()) + .collectAsList(); + Assertions.assertThat(validFiles).hasSize(partitionCount * datafilesPerPartitionCount); + Assertions.assertThat(table.partitionStatisticsFiles()).isEmpty(); + + // -- local compute -------- + long base = System.currentTimeMillis(); + SparkActions actions = SparkActions.get(); + ComputePartitionStats.Result result = + actions.computePartitionStatistics(table).localCompute(true).execute(); + Assertions.assertThat(table.partitionStatisticsFiles()).containsExactly(result.outputFile()); + System.out.println( + "#### time taken for local compute in milli: " + (System.currentTimeMillis() - base)); + + // read the partition entries from the stats file + Schema schema = + PartitionEntry.icebergSchema(Partitioning.partitionType(table.specs().values())); + List rows; + try (CloseableIterable recordIterator = + PartitionStatsUtil.readPartitionStatsFile( + schema, Files.localInput(result.outputFile().path()))) { + rows = Lists.newArrayList(recordIterator); + } catch (IOException e) { + throw new RuntimeException(e); + } + Assertions.assertThat(rows.size()).isEqualTo(partitionCount); + + // ---- distributed compute --- + table + .updatePartitionStatistics() + .removePartitionStatistics(result.outputFile().snapshotId()) + .commit(); + Assertions.assertThat(table.partitionStatisticsFiles()).isEmpty(); + + base = System.currentTimeMillis(); + actions = SparkActions.get(); + result = actions.computePartitionStatistics(table).localCompute(false).execute(); + Assertions.assertThat(table.partitionStatisticsFiles()).containsExactly(result.outputFile()); + System.out.println( + "#### time taken for distributed compute in milli: " + (System.currentTimeMillis() - base)); + + // can't use PartitionStatsUtil.readPartitionStatsFile as it uses + // ParquetAvroValueReaders$ReadBuilder + // and since native parquet doesn't write column ids, reader throws NPE. + List output = + spark + .read() + .parquet(result.outputFile().path()) + .select("PARTITION_DATA", "DATA_RECORD_COUNT", "DATA_FILE_COUNT") + .collectAsList(); + Assertions.assertThat(output.size()).isEqualTo(partitionCount); + } +}