Add per-row resolution of schema for select in BigQueryLoad

Author: qqndrew

API/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <groupId>org.ohnlp.backbone</groupId>
         <artifactId>backbone-parent</artifactId>
-        <version>1.1.2</version>
+        <version>1.1.3</version>
     </parent>
 
     <artifactId>api</artifactId>

Core/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <groupId>org.ohnlp.backbone</groupId>
         <artifactId>backbone-parent</artifactId>
-        <version>1.1.2</version>
+        <version>1.1.3</version>
     </parent>
 
     <artifactId>core</artifactId>

Example-Backbone-Configs/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <groupId>org.ohnlp.backbone</groupId>
         <artifactId>backbone-parent</artifactId>
-        <version>1.1.2</version>
+        <version>1.1.3</version>
     </parent>
 
     <artifactId>example-backbone-configs</artifactId>

IO/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <groupId>org.ohnlp.backbone</groupId>
         <artifactId>backbone-parent</artifactId>
-        <version>1.1.2</version>
+        <version>1.1.3</version>
     </parent>
 
     <groupId>org.ohnlp.backbone.io</groupId>

IO/src/main/java/org/ohnlp/backbone/io/bigquery/BigQueryLoad.java

@@ -4,13 +4,15 @@
 import com.google.api.services.bigquery.model.TableRow;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryIO;
 import org.apache.beam.sdk.io.gcp.bigquery.BigQueryUtils;
+import org.apache.beam.sdk.schemas.Schema;
 import org.apache.beam.sdk.schemas.transforms.Select;
 import org.apache.beam.sdk.transforms.DoFn;
 import org.apache.beam.sdk.transforms.ParDo;
 import org.apache.beam.sdk.values.PCollection;
 import org.apache.beam.sdk.values.PDone;
 import org.apache.beam.sdk.values.POutput;
 import org.apache.beam.sdk.values.Row;
+import org.joda.time.Duration;
 import org.ohnlp.backbone.api.Load;
 import org.ohnlp.backbone.api.exceptions.ComponentInitializationException;
 
@@ -36,13 +38,32 @@ public void initFromConfig(JsonNode config) throws ComponentInitializationExcept
     @Override
     public POutput expand(PCollection<Row> input) {
         if (this.selectedFields.size() > 0) {
-            input = input.apply("Subset Columns", Select.fieldNames(this.selectedFields.toArray(new String[0])));
+            input = input.apply("Subset Columns", ParDo.of(new DoFn<Row, Row>() {
+                @ProcessElement
+                public void process(ProcessContext c) {
+                    Row input = c.element();
+                    // We have to dynamically resolve schema row by row because
+                    // we don't store schema in the collection itself as it is user-configurable/dynamic
+                    Schema inputSchema = input.getSchema();
+                    List<Schema.Field> outputSchemaFields = new ArrayList<>();
+                    for (String fieldName : selectedFields) {
+                        outputSchemaFields.add(inputSchema.getField(fieldName));
+                    }
+                    Schema outSchema = Schema.of(outputSchemaFields.toArray(new Schema.Field[0]));
+                    // And now just map the values
+                    List<Object> outputValues = new ArrayList<>();
+                    for (String s : selectedFields) {
+                        outputValues.add(input.getValue(s));
+                    }
+                    c.output(Row.withSchema(outSchema).addValues(outputValues).build());
+                }
+            }));
         }
         return input.apply(
                 "Transform output rows to BigQuery TableRow format", ParDo.of(
                         new DoFn<Row, TableRow>() {
                             @ProcessElement
-                            public void processELement(@Element Row row, OutputReceiver<TableRow> out) {
+                            public void processElement(@Element Row row, OutputReceiver<TableRow> out) {
                                 out.output(BigQueryUtils.toTableRow(row));
                             }
                         }

Plugin-Manager/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <groupId>org.ohnlp.backbone</groupId>
         <artifactId>backbone-parent</artifactId>
-        <version>1.1.2</version>
+        <version>1.1.3</version>
     </parent>
 
     <artifactId>plugin-manager</artifactId>

Transforms/pom.xml

@@ -7,7 +7,7 @@
     <parent>
         <groupId>org.ohnlp.backbone</groupId>
         <artifactId>backbone-parent</artifactId>
-        <version>1.1.2</version>
+        <version>1.1.3</version>
     </parent>
 
     <groupId>org.ohnlp.backbone.transforms</groupId>

pom.xml

@@ -6,7 +6,7 @@
 
     <groupId>org.ohnlp.backbone</groupId>
     <artifactId>backbone-parent</artifactId>
-    <version>1.1.2</version>
+    <version>1.1.3</version>
 
 
     <properties>