diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java index d06c3d764c4f..d25ce86f60d7 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProto.java @@ -848,6 +848,34 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) return Iterables.getOnlyElement(fileDescriptor.getMessageTypes()); } + private static final List TABLE_ROW_MODEL_FIELDS = + ImmutableList.copyOf(new TableRow().getClassInfo().getNames()); + private static final boolean TABLE_ROW_HAS_ONLY_F = + TABLE_ROW_MODEL_FIELDS.size() == 1 && TABLE_ROW_MODEL_FIELDS.contains("f"); + + private static Map unknownFieldMap(TableRow row, String key) { + // Model fields need GenericData's reflective access and its setter/removal semantics. + return TABLE_ROW_MODEL_FIELDS.contains(key) ? row : row.getUnknownKeys(); + } + + private static void setUnknownField(TableRow row, String key, Object value) { + if (TABLE_ROW_MODEL_FIELDS.contains(key)) { + row.set(key, value); + } else { + row.getUnknownKeys().put(key, value); + } + } + + private static boolean isUnknownMapEmpty(Map map) { + if (TABLE_ROW_HAS_ONLY_F && map instanceof TableRow) { + TableRow row = (TableRow) map; + // A non-null model field counts as an entry, even when its cell list is empty. + return row.getUnknownKeys().isEmpty() && row.getF() == null; + } + // Preserve Map implementations and any future TableRow model fields. + return map.isEmpty(); + } + /** * Given a BigQuery row map, returns a protocol-buffer message for the BigQuery Storage API. * @@ -895,7 +923,7 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) if (fieldDescriptor == null) { if (unknownFields != null) { - unknownFields.set(key, entry.getValue()); + setUnknownField(unknownFields, key, entry.getValue()); } if (ignoreUnknownValues) { continue; @@ -978,11 +1006,13 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) TableRow nestedUnknown = new TableRow(); if (fieldDescriptor.isRepeated()) { ((List) - unknownFields.computeIfAbsent(key, k -> new ArrayList())) + unknownFieldMap(unknownFields, key) + .computeIfAbsent(key, k -> new ArrayList())) .add(nestedUnknown); return nestedUnknown; } - return (TableRow) unknownFields.computeIfAbsent(key, k -> nestedUnknown); + return (TableRow) + unknownFieldMap(unknownFields, key).computeIfAbsent(key, k -> nestedUnknown); }; @Nullable Object value = @@ -1000,17 +1030,19 @@ public static Descriptor wrapDescriptorProto(DescriptorProto descriptorProto) // For STRUCT fields, we add a placeholder to unknownFields using the getNestedUnknown // supplier (in case we encounter unknown nested fields). If the placeholder comes out // to be empty, we should clean it up - if ((fieldSchemaInformation.getType().equals(TableFieldSchema.Type.STRUCT) - && unknownFields != null) - && ((unknownFields.get(key) instanceof Map - && ((Map) unknownFields.get(key)).isEmpty()) // single struct, empty - || (unknownFields.get(key) - instanceof List // repeated struct, empty list or list with empty structs - && (((List) unknownFields.get(key)).isEmpty() - || ((List) unknownFields.get(key)) - .stream() - .allMatch(row -> row == null || ((Map) row).isEmpty()))))) { - unknownFields.remove(key); + if (fieldSchemaInformation.getType().equals(TableFieldSchema.Type.STRUCT) + && unknownFields != null) { + Map parentUnknown = unknownFieldMap(unknownFields, key); + Object nestedUnknown = parentUnknown.get(key); + if ((nestedUnknown instanceof Map && isUnknownMapEmpty((Map) nestedUnknown)) + || (nestedUnknown instanceof List + && (((List) nestedUnknown).isEmpty() + || ((List) nestedUnknown) + .stream() + .allMatch( + row -> row == null || isUnknownMapEmpty((Map) row))))) { + parentUnknown.remove(key); + } } } catch (Exception e) { throw new SchemaDoesntMatchException( diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoUnknownFieldsTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoUnknownFieldsTest.java new file mode 100644 index 000000000000..e4b1751850c6 --- /dev/null +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/TableRowToStorageApiProtoUnknownFieldsTest.java @@ -0,0 +1,301 @@ +/* + * 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.beam.sdk.io.gcp.bigquery; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertSame; +import static org.junit.Assert.assertThrows; +import static org.junit.Assert.assertTrue; + +import com.google.api.services.bigquery.model.TableCell; +import com.google.api.services.bigquery.model.TableFieldSchema; +import com.google.api.services.bigquery.model.TableRow; +import com.google.api.services.bigquery.model.TableSchema; +import com.google.protobuf.Descriptors.Descriptor; +import com.google.protobuf.DynamicMessage; +import java.io.ByteArrayOutputStream; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; +import java.util.LinkedHashMap; +import java.util.List; +import org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.ErrorCollector; +import org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.SchemaDoesntMatchException; +import org.apache.beam.sdk.io.gcp.bigquery.TableRowToStorageApiProto.SchemaInformation; +import org.junit.Test; +import org.junit.runner.RunWith; +import org.junit.runners.JUnit4; + +/** Regression tests for unknown-column collector attachment and pruning. */ +@RunWith(JUnit4.class) +@SuppressWarnings("nullness") // TODO(https://github.com/apache/beam/issues/20497) +public class TableRowToStorageApiProtoUnknownFieldsTest { + private static final TableFieldSchema KNOWN_FIELD = + new TableFieldSchema().setName("known").setType("STRING"); + private static final TableSchema SCHEMA = + new TableSchema() + .setFields( + Arrays.asList( + new TableFieldSchema() + .setName("record") + .setType("RECORD") + .setFields(Collections.singletonList(KNOWN_FIELD)), + new TableFieldSchema() + .setName("records") + .setType("RECORD") + .setMode("REPEATED") + .setFields(Collections.singletonList(KNOWN_FIELD)))); + + private static LinkedHashMap map(Object... entries) { + LinkedHashMap result = new LinkedHashMap<>(); + for (int i = 0; i < entries.length; i += 2) { + result.put((String) entries[i], entries[i + 1]); + } + return result; + } + + private static DynamicMessage convert(AbstractMap row, TableRow unknown) + throws Exception { + return convert(SCHEMA, row, unknown); + } + + private static DynamicMessage convert( + TableSchema schema, AbstractMap row, TableRow unknown) throws Exception { + return TableRowToStorageApiProto.messageFromMap( + SchemaInformation.fromTableSchema(schema), + TableRowToStorageApiProto.getDescriptorFromTableSchema(schema, true, false), + row, + true, + true, + unknown, + null, + null, + ErrorCollector.DONT_COLLECT); + } + + private static byte[] encoded(TableRow row) throws Exception { + ByteArrayOutputStream bytes = new ByteArrayOutputStream(); + TableRowJsonCoder.of().encode(row, bytes); + return bytes.toByteArray(); + } + + private static TableSchema fSchema() { + return new TableSchema() + .setFields( + Collections.singletonList( + new TableFieldSchema() + .setName("f") + .setType("RECORD") + .setFields(Collections.singletonList(KNOWN_FIELD)))); + } + + @Test + public void testKnownSingleStructPrunesPlaceholder() throws Exception { + TableRow unknown = new TableRow(); + DynamicMessage message = convert(map("record", map("known", "value")), unknown); + Descriptor descriptor = message.getDescriptorForType(); + assertTrue(unknown.isEmpty()); + assertTrue(message.getField(descriptor.findFieldByName("record")).toString().contains("value")); + } + + @Test + public void testKnownRepeatedStructPrunesPlaceholders() throws Exception { + TableRow unknown = new TableRow(); + DynamicMessage message = + convert(map("records", Arrays.asList(map("known", "a"), map("known", "b"))), unknown); + Descriptor descriptor = message.getDescriptorForType(); + assertTrue(unknown.isEmpty()); + assertEquals(2, message.getRepeatedFieldCount(descriptor.findFieldByName("records"))); + } + + @Test + public void testEmptyRepeatedStructPrunesPlaceholder() throws Exception { + TableRow unknown = new TableRow(); + convert(map("records", Collections.emptyList()), unknown); + assertTrue(unknown.isEmpty()); + } + + @Test + public void testSingleStructPreservesUnknownJsonAndProtoBytes() throws Exception { + TableRow unknown = new TableRow(); + DynamicMessage actual = convert(map("record", map("known", "value", "added", "new")), unknown); + DynamicMessage expected = convert(map("record", map("known", "value")), new TableRow()); + assertArrayEquals(expected.toByteArray(), actual.toByteArray()); + assertArrayEquals( + encoded(new TableRow().set("record", new TableRow().set("added", "new"))), + encoded(unknown)); + } + + @Test + public void testRepeatedStructPreservesUnknownPositionsAndProtoBytes() throws Exception { + for (int position = 0; position < 3; position++) { + List rows = new ArrayList<>(); + List knownRows = new ArrayList<>(); + List expectedUnknown = new ArrayList<>(); + for (int index = 0; index < 3; index++) { + knownRows.add(map("known", "value-" + index)); + rows.add( + index == position + ? map("known", "value-" + index, "added", "new") + : map("known", "value-" + index)); + expectedUnknown.add( + index == position ? new TableRow().set("added", "new") : new TableRow()); + } + TableRow unknown = new TableRow(); + DynamicMessage actual = convert(map("records", rows), unknown); + DynamicMessage expected = convert(map("records", knownRows), new TableRow()); + assertArrayEquals(expected.toByteArray(), actual.toByteArray()); + assertArrayEquals(encoded(new TableRow().set("records", expectedUnknown)), encoded(unknown)); + } + } + + @Test + public void testRetainsPreexistingUnknownStruct() throws Exception { + TableRow nested = new TableRow().set("old", "value"); + TableRow unknown = new TableRow().set("record", nested); + convert(map("record", map()), unknown); + assertSame(nested, unknown.get("record")); + assertArrayEquals( + encoded(new TableRow().set("record", new TableRow().set("old", "value"))), + encoded(unknown)); + } + + @Test + public void testPrunesPreexistingNullAndEmptyRepeatedPlaceholders() throws Exception { + TableRow unknown = new TableRow().set("records", Arrays.asList(null, new TableRow())); + convert(map("records", Collections.emptyList()), unknown); + assertTrue(unknown.isEmpty()); + } + + @Test + public void testNonNullModelCellListKeepsSingleStructNonempty() throws Exception { + for (List cells : + Arrays.asList( + Collections.emptyList(), + Collections.singletonList(new TableCell().setV("cell")))) { + TableRow nested = new TableRow().setF(cells); + TableRow unknown = new TableRow().set("record", nested); + convert(map("record", map("known", "value")), unknown); + assertSame(nested, unknown.get("record")); + assertSame(cells, nested.getF()); + assertArrayEquals( + encoded(new TableRow().set("record", new TableRow().setF(cells))), encoded(unknown)); + } + } + + @Test + public void testEmptyModelCellListKeepsRepeatedPlaceholderNonempty() throws Exception { + TableRow nested = new TableRow().setF(Collections.emptyList()); + List rows = Arrays.asList(null, nested, new TableRow()); + TableRow unknown = new TableRow().set("records", rows); + convert(map("records", Collections.emptyList()), unknown); + assertSame(rows, unknown.get("records")); + } + + @Test + public void testOrdinaryMapsKeepTheirEmptinessBehavior() throws Exception { + LinkedHashMap nested = map("old", "value"); + TableRow unknown = new TableRow().set("record", nested); + convert(map("record", null), unknown); + assertSame(nested, unknown.get("record")); + TableRow empty = new TableRow().set("record", map()); + convert(map("record", null), empty); + assertTrue(empty.isEmpty()); + } + + @Test + public void testUnknownFUsesModelSetter() throws Exception { + List cells = Collections.singletonList(new TableCell().setV("value")); + TableRow unknown = new TableRow(); + convert(map("f", cells), unknown); + assertSame(cells, unknown.getF()); + assertFalse(unknown.getUnknownKeys().containsKey("f")); + } + + @Test + public void testUnknownFPreservesShadowedBackingKey() throws Exception { + List cells = Collections.singletonList(new TableCell().setV("value")); + TableRow unknown = new TableRow(); + unknown.getUnknownKeys().put("f", "shadow"); + convert(map("f", cells), unknown); + assertSame(cells, unknown.getF()); + assertEquals("shadow", unknown.getUnknownKeys().get("f")); + } + + @Test + public void testUnknownFPreservesSetterTypeError() { + TableRow unknown = new TableRow(); + assertThrows( + IllegalArgumentException.class, () -> convert(map("f", map("known", "value")), unknown)); + assertFalse(unknown.getUnknownKeys().containsKey("f")); + } + + @Test + public void testKnownFPreservesAttachmentTypeError() { + SchemaDoesntMatchException error = + assertThrows( + SchemaDoesntMatchException.class, + () -> convert(fSchema(), map("f", map("known", "value")), new TableRow())); + assertTrue(error.getCause() instanceof IllegalArgumentException); + } + + @Test + public void testKnownFPreservesUnsupportedModelRemoval() { + TableRow unknown = new TableRow().setF(Collections.emptyList()); + SchemaDoesntMatchException error = + assertThrows( + SchemaDoesntMatchException.class, () -> convert(fSchema(), map("f", null), unknown)); + assertTrue(error.getCause() instanceof UnsupportedOperationException); + assertEquals(Collections.emptyList(), unknown.getF()); + } + + @Test + public void testNullCollectorPreservesProtoBytes() throws Exception { + AbstractMap row = map("record", map("known", "value", "added", "new")); + assertArrayEquals(convert(row, new TableRow()).toByteArray(), convert(row, null).toByteArray()); + } + + @Test + public void testCellFormatUsesPublicConverterPath() throws Exception { + TableRow row = + new TableRow() + .setF( + Arrays.asList( + new TableCell().setV(map("known", "value")), + new TableCell().setV(Collections.emptyList()))); + DynamicMessage actual = + TableRowToStorageApiProto.messageFromTableRow( + SchemaInformation.fromTableSchema(SCHEMA), + TableRowToStorageApiProto.getDescriptorFromTableSchema(SCHEMA, true, false), + row, + true, + true, + new TableRow(), + null, + -1, + ErrorCollector.DONT_COLLECT); + DynamicMessage expected = + convert( + map("record", map("known", "value"), "records", Collections.emptyList()), + new TableRow()); + assertArrayEquals(expected.toByteArray(), actual.toByteArray()); + } +}