From ec3897b65a699ad55ca99c96ac3f46cb0682b6fd Mon Sep 17 00:00:00 2001 From: Shivam Malhotra Date: Tue, 10 Dec 2024 14:40:53 -0600 Subject: [PATCH] Review with Devin --- .../IcebergKeyValuePartitionedLayout.java | 41 ++++++++++--------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergKeyValuePartitionedLayout.java b/extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergKeyValuePartitionedLayout.java index 0b727a17c70..82522d471a0 100644 --- a/extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergKeyValuePartitionedLayout.java +++ b/extensions/iceberg/src/main/java/io/deephaven/iceberg/layout/IcebergKeyValuePartitionedLayout.java @@ -17,8 +17,6 @@ import java.net.URI; import java.util.*; -import java.util.concurrent.atomic.AtomicInteger; -import java.util.stream.Collectors; /** * Iceberg {@link TableLocationKeyFinder location finder} for tables with partitions that will discover data files from @@ -53,24 +51,27 @@ public IcebergKeyValuePartitionedLayout( // We can assume due to upstream validation that there are no duplicate names (after renaming) that are included // in the output definition, so we can ignore duplicates. - final AtomicInteger icebergIndex = new AtomicInteger(0); - // TODO (DH-18160): Improve support for handling non-identity transforms - identityPartitioningColumns = partitionSpec.fields().stream() - .filter(partitionField -> partitionField.transform().isIdentity()) - .map(PartitionField::name) - .map(icebergColName -> { - final String dhColName = instructions.columnRenames().getOrDefault(icebergColName, icebergColName); - final ColumnDefinition columnDef = tableDef.getColumn(dhColName); - if (tableDef.getColumn(dhColName) == null) { - throw new TableDataException("Partitioning column " + dhColName + " not found in table " + - "definition but corresponding identity partitioning column " + icebergColName + " is " + - "present in the partition spec, table definition: " + tableDef + ", partition spec: " + - partitionSpec); - } - return new IdentityPartitioningColData(dhColName, TypeUtils.getBoxedType(columnDef.getDataType()), - icebergIndex.getAndIncrement()); - }) - .collect(Collectors.toList()); + final List partitionFields = partitionSpec.fields(); + final int numPartitionFields = partitionFields.size(); + identityPartitioningColumns = new ArrayList<>(numPartitionFields); + for (int fieldId = 0; fieldId < numPartitionFields; ++fieldId) { + final PartitionField partitionField = partitionFields.get(fieldId); + if (!partitionField.transform().isIdentity()) { + // TODO (DH-18160): Improve support for handling non-identity transforms + continue; + } + final String icebergColName = partitionField.name(); + final String dhColName = instructions.columnRenames().getOrDefault(icebergColName, icebergColName); + final ColumnDefinition columnDef = tableDef.getColumn(dhColName); + if (columnDef == null) { + throw new TableDataException("Partitioning column " + dhColName + " not found in table definition " + + "but corresponding identity partitioning column " + icebergColName + " is present in the " + + "partition spec, table definition: " + tableDef + ", partition spec: " + partitionSpec); + } + identityPartitioningColumns.add(new IdentityPartitioningColData(dhColName, + TypeUtils.getBoxedType(columnDef.getDataType()), fieldId)); + + } } @Override