[IcebergIO] Fix conversion logic for arrays of structs and maps of structs; fix output Schema resolution with column pruning

.github/trigger_files/IO_Iceberg_Integration_Tests.json

@@ -1,5 +1,4 @@
 {
  "comment": "Modify this file in a trivial way to cause this test suite to run.",
- "modification": 3,
- "https://github.com/apache/beam/pull/35159": "moving WindowedValue and making an interface"
+ "modification": 5
 }

CHANGES.md

@@ -75,6 +75,8 @@
 ## I/Os
 
 * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)).
+* [IcebergIO] Support reading with column pruning ([#34856](https://github.com/apache/beam/pull/34856))
+* [IcebergIO] Support reading with pushdown filtering ([#34827](https://github.com/apache/beam/pull/34827))
 
 ## New Features / Improvements
 * Adding Google Storage Requests Pays feature (Golang)([#30747](https://github.com/apache/beam/issues/30747)).
@@ -103,6 +105,7 @@
 * [Python] Fixed vLLM leaks connections causing a throughput bottleneck and underutilization of GPU ([35053](https://github.com/apache/beam/pull/35053))
 * (Python) Fixed cloudpickle overwriting class states every time loading a same object of dynamic class ([#35062](https://github.com/apache/beam/issues/35062)).
 * [Python] Fixed pip install apache-beam[interactive] causes crash on google colab ([#35148](https://github.com/apache/beam/pull/35148)).
+* [IcebergIO] Fixed Beam <-> Iceberg conversion logic for arrays of structs and maps of structs ([#35230](https://github.com/apache/beam/pull/35230)).
 
 ## Security Fixes
 * Fixed [CVE-YYYY-NNNN](https://www.cve.org/CVERecord?id=CVE-YYYY-NNNN) (Java/Python/Go) ([#X](https://github.com/apache/beam/issues/X)).

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/FilterUtils.java

@@ -24,9 +24,11 @@
 import java.time.LocalDate;
 import java.time.LocalDateTime;
 import java.time.LocalTime;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
 import java.util.function.BiFunction;
 import java.util.stream.Collectors;
 import org.apache.beam.vendor.calcite.v1_28_0.org.apache.calcite.sql.SqlBasicCall;
@@ -72,6 +74,55 @@ class FilterUtils {
  .put(SqlKind.OR, Operation.OR)
  .build();
 
+ /**
+ * Parses a SQL filter expression string and returns a set of all field names referenced within
+ * it.
+ */
+ static Set<String> getReferencedFieldNames(@Nullable String filter) {
+ if (filter == null || filter.trim().isEmpty()) {
+ return new HashSet<>();
+ }
+
+ SqlParser parser = SqlParser.create(filter);
+ try {
+ SqlNode expression = parser.parseExpression();
+ Set<String> fieldNames = new HashSet<>();
+ extractFieldNames(expression, fieldNames);
+ return fieldNames;
+ } catch (Exception exception) {
+ throw new RuntimeException(
+ String.format("Encountered an error when parsing filter: '%s'", filter), exception);
+ }
+ }
+
+ private static void extractFieldNames(SqlNode node, Set<String> fieldNames) {
+ if (node instanceof SqlIdentifier) {
+ fieldNames.add(((SqlIdentifier) node).getSimple());
+ } else if (node instanceof SqlBasicCall) {
+ // recursively check operands
+ SqlBasicCall call = (SqlBasicCall) node;
+ for (SqlNode operand : call.getOperandList()) {
+ extractFieldNames(operand, fieldNames);
+ }
+ } else if (node instanceof SqlNodeList) {
+ // For IN clauses, the right-hand side is a SqlNodeList, so iterate through its elements
+ SqlNodeList nodeList = (SqlNodeList) node;
+ for (SqlNode element : nodeList.getList()) {
+ if (element != null) {
+ extractFieldNames(element, fieldNames);
+ }
+ }
+ }
+ // SqlLiteral nodes do not contain field names, so we can ignore them.
+ }
+
+ /**
+ * parses a SQL filter expression string into an Iceberg {@link Expression} that can be used for
+ * data pruning.
+ *
+ * <p>Note: This utility currently supports only top-level fields within the filter expression.
+ * Nested field references are not supported.
+ */
  static Expression convert(@Nullable String filter, Schema schema) {
  if (filter == null) {
  return Expressions.alwaysTrue();

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergIO.java

@@ -603,7 +603,7 @@ public PCollection<Row> expand(PBegin input) {
  .setUseCdc(getUseCdc())
  .setKeepFields(getKeep())
  .setDropFields(getDrop())
- .setFilter(FilterUtils.convert(getFilter(), table.schema()))
+ .setFilterString(getFilter())
  .build();
  scanConfig.validate(table);
 

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergScanConfig.java

@@ -31,6 +31,7 @@
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.MoreObjects;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.Table;
 import org.apache.iceberg.catalog.TableIdentifier;
@@ -45,6 +46,10 @@
 @AutoValue
 public abstract class IcebergScanConfig implements Serializable {
  private transient @MonotonicNonNull Table cachedTable;
+ private transient org.apache.iceberg.@MonotonicNonNull Schema cachedProjectedSchema;
+ private transient org.apache.iceberg.@MonotonicNonNull Schema cachedRequiredSchema;
+ private transient @MonotonicNonNull Evaluator cachedEvaluator;
+ private transient @MonotonicNonNull Expression cachedFilter;
 
  public enum ScanType {
  TABLE,
@@ -75,19 +80,59 @@ public Table getTable() {
  @VisibleForTesting
  static org.apache.iceberg.Schema resolveSchema(
  org.apache.iceberg.Schema schema, @Nullable List<String> keep, @Nullable List<String> drop) {
+ return resolveSchema(schema, keep, drop, null);
+ }
+
+ @VisibleForTesting
+ static org.apache.iceberg.Schema resolveSchema(
+ org.apache.iceberg.Schema schema,
+ @Nullable List<String> keep,
+ @Nullable List<String> drop,
+ @Nullable Set<String> fieldsInFilter) {
+ ImmutableList.Builder<String> selectedFieldsBuilder = ImmutableList.builder();
  if (keep != null && !keep.isEmpty()) {
-schema = schema.select(keep);
+selectedFieldsBuilder.addAll(keep);
  } else if (drop != null && !drop.isEmpty()) {
  Set<String> fields =
  schema.columns().stream().map(Types.NestedField::name).collect(Collectors.toSet());
  drop.forEach(fields::remove);
- schema = schema.select(fields);
+ selectedFieldsBuilder.addAll(fields);
+ } else {
+ // default: include all columns
+ return schema;
  }
- return schema;
+
+ if (fieldsInFilter != null && !fieldsInFilter.isEmpty()) {
+ fieldsInFilter.stream()
+ .map(f -> schema.caseInsensitiveFindField(f).name())
+ .forEach(selectedFieldsBuilder::add);
+ }
+ ImmutableList<String> selectedFields = selectedFieldsBuilder.build();
+ return selectedFields.isEmpty() ? schema : schema.select(selectedFields);
  }
 
+ /** Returns the projected Schema after applying column pruning. */
  public org.apache.iceberg.Schema getProjectedSchema() {
- return resolveSchema(getTable().schema(), getKeepFields(), getDropFields());
+ if (cachedProjectedSchema == null) {
+ cachedProjectedSchema = resolveSchema(getTable().schema(), getKeepFields(), getDropFields());
+ }
+ return cachedProjectedSchema;
+ }
+
+ /**
+ * Returns a Schema that includes all the fields required for a successful read. This includes
+ * explicitly selected fields and fields referenced in the filter statement.
+ */
+ public org.apache.iceberg.Schema getRequiredSchema() {
+ if (cachedRequiredSchema == null) {
+ cachedRequiredSchema =
+ resolveSchema(
+ getTable().schema(),
+ getKeepFields(),
+ getDropFields(),
+ FilterUtils.getReferencedFieldNames(getFilterString()));
+ }
+ return cachedRequiredSchema;
  }
 
  @Pure
@@ -98,15 +143,22 @@ public Evaluator getEvaluator() {
  return null;
  }
  if (cachedEvaluator == null) {
- cachedEvaluator = new Evaluator(getProjectedSchema().asStruct(), filter);
+ cachedEvaluator = new Evaluator(getRequiredSchema().asStruct(), filter);
  }
  return cachedEvaluator;
  }
 
- private transient @Nullable Evaluator cachedEvaluator;
+ @Pure
+ @Nullable
+ public Expression getFilter() {
+ if (cachedFilter == null) {
+ cachedFilter = FilterUtils.convert(getFilterString(), getTable().schema());
+ }
+ return cachedFilter;
+ }
 
  @Pure
- public abstract @Nullable Expression getFilter();
+ public abstract @Nullable String getFilterString();
 
  @Pure
  public abstract @Nullable Boolean getCaseSensitive();
@@ -172,7 +224,7 @@ public Evaluator getEvaluator() {
  public static Builder builder() {
  return new AutoValue_IcebergScanConfig.Builder()
  .setScanType(ScanType.TABLE)
- .setFilter(null)
+ .setFilterString(null)
  .setCaseSensitive(null)
  .setOptions(ImmutableMap.of())
  .setSnapshot(null)
@@ -211,7 +263,7 @@ public Builder setTableIdentifier(String... names) {
 
  public abstract Builder setSchema(Schema schema);
 
- public abstract Builder setFilter(@Nullable Expression filter);
+ public abstract Builder setFilterString(@Nullable String filter);
 
  public abstract Builder setCaseSensitive(@Nullable Boolean caseSensitive);
 

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/IcebergUtils.java

@@ -18,6 +18,8 @@
 package org.apache.beam.sdk.io.iceberg;
 
 import static org.apache.beam.sdk.util.Preconditions.checkArgumentNotNull;
+import static org.apache.beam.sdk.util.Preconditions.checkStateNotNull;
+import static org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.base.Preconditions.checkState;
 
 import java.nio.ByteBuffer;
 import java.time.LocalDate;
@@ -30,17 +32,20 @@
 import java.util.Map;
 import java.util.Optional;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.logicaltypes.SqlTypes;
 import org.apache.beam.sdk.util.Preconditions;
 import org.apache.beam.sdk.values.Row;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.annotations.VisibleForTesting;
+import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableList;
 import org.apache.beam.vendor.guava.v32_1_2_jre.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.data.GenericRecord;
 import org.apache.iceberg.data.Record;
 import org.apache.iceberg.types.Type;
 import org.apache.iceberg.types.Types;
 import org.apache.iceberg.util.DateTimeUtil;
+import org.checkerframework.checker.nullness.qual.NonNull;
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.joda.time.DateTime;
 import org.joda.time.Instant;
@@ -345,10 +350,34 @@ private static void copyFieldIntoRecord(Record rec, Types.NestedField field, Row
  copyRowIntoRecord(GenericRecord.create(field.type().asStructType()), row)));
  break;
  case LIST:
- Optional.ofNullable(value.getArray(name)).ifPresent(list -> rec.setField(name, list));
+ Iterable<@NonNull ?> icebergList = value.getIterable(name);
+ Type collectionType = ((Types.ListType) field.type()).elementType();
+
+ if (collectionType.isStructType() && icebergList != null) {
+ org.apache.iceberg.Schema innerSchema = collectionType.asStructType().asSchema();
+ ImmutableList.Builder<Record> builder = ImmutableList.builder();
+ for (Row v : (Iterable<Row>) icebergList) {
+ builder.add(beamRowToIcebergRecord(innerSchema, v));
+ }
+ icebergList = builder.build();
+ }
+ Optional.ofNullable(icebergList).ifPresent(list -> rec.setField(name, list));
  break;
  case MAP:
- Optional.ofNullable(value.getMap(name)).ifPresent(v -> rec.setField(name, v));
+ Map<?, ?> icebergMap = value.getMap(name);
+ Type valueType = ((Types.MapType) field.type()).valueType();
+ // recurse on struct types
+ if (valueType.isStructType() && icebergMap != null) {
+ org.apache.iceberg.Schema innerSchema = valueType.asStructType().asSchema();
+
+ ImmutableMap.Builder<Object, Record> newMap = ImmutableMap.builder();
+ for (Map.Entry<?, ?> entry : icebergMap.entrySet()) {
+ Row row = checkStateNotNull(((Row) entry.getValue()));
+ newMap.put(checkStateNotNull(entry.getKey()), beamRowToIcebergRecord(innerSchema, row));
+ }
+ icebergMap = newMap.build();
+ }
+ Optional.ofNullable(icebergMap).ifPresent(v -> rec.setField(name, v));
  break;
  }
  }
@@ -426,10 +455,68 @@ public static Row icebergRecordToBeamRow(Schema schema, Record record) {
  case DOUBLE: // Iceberg and Beam both use double
  case STRING: // Iceberg and Beam both use String
  case BOOLEAN: // Iceberg and Beam both use boolean
+ rowBuilder.addValue(icebergValue);
+ break;
  case ARRAY:
+ checkState(
+ icebergValue instanceof List,
+ "Expected List type for field '%s' but received %s",
+ field.getName(),
+ icebergValue.getClass());
+ List<@NonNull ?> beamList = (List<@NonNull ?>) icebergValue;
+ Schema.FieldType collectionType =
+ checkStateNotNull(field.getType().getCollectionElementType());
+ // recurse on struct types
+ if (collectionType.getTypeName().isCompositeType()) {
+ Schema innerSchema = checkStateNotNull(collectionType.getRowSchema());
+ beamList =
+ beamList.stream()
+ .map(v -> icebergRecordToBeamRow(innerSchema, (Record) v))
+ .collect(Collectors.toList());
+ }
+ rowBuilder.addValue(beamList);
+ break;
  case ITERABLE:
+ checkState(
+ icebergValue instanceof Iterable,
+ "Expected Iterable type for field '%s' but received %s",
+ field.getName(),
+ icebergValue.getClass());
+ Iterable<@NonNull ?> beamIterable = (Iterable<@NonNull ?>) icebergValue;
+ Schema.FieldType iterableCollectionType =
+ checkStateNotNull(field.getType().getCollectionElementType());
+ // recurse on struct types
+ if (iterableCollectionType.getTypeName().isCompositeType()) {
+ Schema innerSchema = checkStateNotNull(iterableCollectionType.getRowSchema());
+ ImmutableList.Builder<Row> builder = ImmutableList.builder();
+ for (Record v : (Iterable<@NonNull Record>) icebergValue) {
+ builder.add(icebergRecordToBeamRow(innerSchema, v));
+ }
+ beamIterable = builder.build();
+ }
+ rowBuilder.addValue(beamIterable);
+ break;
  case MAP:
- rowBuilder.addValue(icebergValue);
+ checkState(
+ icebergValue instanceof Map,
+ "Expected Map type for field '%s' but received %s",
+ field.getName(),
+ icebergValue.getClass());
+ Map<?, ?> beamMap = (Map<?, ?>) icebergValue;
+ Schema.FieldType valueType = checkStateNotNull(field.getType().getMapValueType());
+ // recurse on struct types
+ if (valueType.getTypeName().isCompositeType()) {
+ Schema innerSchema = checkStateNotNull(valueType.getRowSchema());
+ ImmutableMap.Builder<Object, Row> newMap = ImmutableMap.builder();
+ for (Map.Entry<?, ?> entry : ((Map<?, ?>) icebergValue).entrySet()) {
+ Record rec = ((Record) entry.getValue());
+ newMap.put(
+ checkStateNotNull(entry.getKey()),
+ icebergRecordToBeamRow(innerSchema, checkStateNotNull(rec)));
+ }
+ beamMap = newMap.build();
+ }
+ rowBuilder.addValue(beamMap);
  break;
  case DATETIME:
  // Iceberg uses a long for micros.

sdks/java/io/iceberg/src/main/java/org/apache/beam/sdk/io/iceberg/ReadFromTasks.java

@@ -74,10 +74,9 @@ public void process(
  return;
  }
  FileScanTask task = fileScanTasks.get((int) l);
- org.apache.iceberg.Schema projected = scanConfig.getProjectedSchema();
- Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(projected);
+ Schema beamSchema = IcebergUtils.icebergSchemaToBeamSchema(scanConfig.getProjectedSchema());
  try (CloseableIterable<Record> fullIterable =
- ReadUtils.createReader(task, table, projected)) {
+ ReadUtils.createReader(task, table, scanConfig.getRequiredSchema())) {
  CloseableIterable<Record> reader = ReadUtils.maybeApplyFilter(fullIterable, scanConfig);
 
  for (Record record : reader) {