Merge pull request #1137 from data-integrations/bq-direct-source_final

[PLUGIN-871] Adding support for BQ source pushdown in BQ Plugin

Author: sahusanket

src/main/java/io/cdap/plugin/gcp/bigquery/source/BigQuerySource.java

@@ -29,6 +29,8 @@
 import com.google.cloud.kms.v1.CryptoKeyName;
 import com.google.cloud.storage.Storage;
 import com.google.common.base.Strings;
+import com.google.common.collect.ImmutableMap;
+import com.google.gson.Gson;
 import io.cdap.cdap.api.annotation.Description;
 import io.cdap.cdap.api.annotation.Metadata;
 import io.cdap.cdap.api.annotation.MetadataProperty;
@@ -46,9 +48,13 @@
 import io.cdap.cdap.etl.api.batch.BatchSource;
 import io.cdap.cdap.etl.api.batch.BatchSourceContext;
 import io.cdap.cdap.etl.api.connector.Connector;
+import io.cdap.cdap.etl.api.engine.sql.SQLEngineInput;
 import io.cdap.cdap.etl.api.validation.ValidationFailure;
 import io.cdap.plugin.common.LineageRecorder;
 import io.cdap.plugin.gcp.bigquery.connector.BigQueryConnector;
+import io.cdap.plugin.gcp.bigquery.sqlengine.BigQueryReadDataset;
+import io.cdap.plugin.gcp.bigquery.sqlengine.BigQuerySQLEngine;
+import io.cdap.plugin.gcp.bigquery.sqlengine.BigQueryWrite;
 import io.cdap.plugin.gcp.bigquery.util.BigQueryConstants;
 import io.cdap.plugin.gcp.bigquery.util.BigQueryUtil;
 import io.cdap.plugin.gcp.common.CmekUtils;
@@ -61,6 +67,7 @@
 
 import java.time.DateTimeException;
 import java.time.LocalDate;
+import java.util.List;
 import java.util.UUID;
 import java.util.stream.Collectors;
 import javax.annotation.Nullable;
@@ -76,6 +83,7 @@
 @Metadata(properties = {@MetadataProperty(key = Connector.PLUGIN_TYPE, value = BigQueryConnector.NAME)})
 public final class BigQuerySource extends BatchSource<LongWritable, GenericData.Record, StructuredRecord> {
   private static final Logger LOG = LoggerFactory.getLogger(BigQuerySource.class);
+  private static final Gson GSON = new Gson();
   public static final String NAME = "BigQueryTable";
   private BigQuerySourceConfig config;
   private Schema outputSchema;
@@ -165,7 +173,7 @@ public void prepareRun(BatchSourceContext context) throws Exception {
     // We call emitLineage before since it creates the dataset with schema.
     Type sourceTableType = config.getSourceTableType();
     emitLineage(context, configuredSchema, sourceTableType, config.getTable());
-    setInputFormat(context);
+    setInputFormat(context, configuredSchema);
   }
 
   @Override
@@ -335,8 +343,31 @@ private void validatePartitionProperties(FailureCollector collector) {
     }
   }
 
-  private void setInputFormat(BatchSourceContext context) {
+  private void setInputFormat(BatchSourceContext context,
+                              Schema configuredSchema) {
+    // Set input for Spark
     context.setInput(Input.of(config.referenceName, new BigQueryInputFormatProvider(configuration)));
+
+    // Add output for SQL Engine Direct read
+    ImmutableMap.Builder<String, String> arguments = new ImmutableMap.Builder<>();
+
+    if (configuredSchema == null) {
+      LOG.debug("BigQuery SQL Engine Input was not initialized. Schema was empty.");
+      return;
+    }
+
+    List<String> fieldNames = configuredSchema.getFields().stream().map(f -> f.getName()).collect(Collectors.toList());
+
+    arguments
+      .put(BigQueryReadDataset.SQL_INPUT_CONFIG, GSON.toJson(config))
+      .put(BigQueryReadDataset.SQL_INPUT_SCHEMA, GSON.toJson(configuredSchema))
+      .put(BigQueryReadDataset.SQL_INPUT_FIELDS, GSON.toJson(fieldNames));
+
+    Input sqlEngineInput = new SQLEngineInput(config.referenceName,
+                                              context.getStageName(),
+                                              BigQuerySQLEngine.class.getName(),
+                                              arguments.build());
+    context.setInput(sqlEngineInput);
   }
 
   private void emitLineage(BatchSourceContext context, Schema schema, Type sourceTableType,

src/main/java/io/cdap/plugin/gcp/bigquery/source/PartitionedBigQueryInputFormat.java

@@ -48,7 +48,6 @@
  * in order to create input splits.
  */
 public class PartitionedBigQueryInputFormat extends AbstractBigQueryInputFormat<LongWritable, GenericData.Record> {
-  private static final String DEFAULT_COLUMN_NAME = "_PARTITIONTIME";
 
   private InputFormat<LongWritable, GenericData.Record> delegateInputFormat =
     new AvroBigQueryInputFormat();
@@ -160,8 +159,8 @@ String generateQuery(String partitionFromDate, String partitionToDate, String fi
     StringBuilder condition = new StringBuilder();
 
     if (timePartitioning != null) {
-      String timePartitionCondition = generateTimePartitionCondition(tableDefinition, timePartitioning,
-                                                                     partitionFromDate, partitionToDate);
+      String timePartitionCondition = BigQueryUtil.generateTimePartitionCondition(tableDefinition, partitionFromDate,
+                                                                                  partitionToDate);
       condition.append(timePartitionCondition);
     }
 
@@ -289,33 +288,4 @@ private static JobReference getJobReference(Configuration conf, BigQueryHelper b
     }
     return new JobReference().setProjectId(projectId).setJobId(savedJobId).setLocation(location);
   }
-
-  private String generateTimePartitionCondition(StandardTableDefinition tableDefinition,
-                                                TimePartitioning timePartitioning, String partitionFromDate,
-                                                String partitionToDate) {
-    StringBuilder timePartitionCondition = new StringBuilder();
-    String columnName = timePartitioning.getField() != null ? timePartitioning.getField() : DEFAULT_COLUMN_NAME;
-
-    LegacySQLTypeName columnType = null;
-    if (!DEFAULT_COLUMN_NAME.equals(columnName)) {
-      columnType = tableDefinition.getSchema().getFields().get(columnName).getType();
-    }
-
-    String columnNameTS = columnName;
-    if (!LegacySQLTypeName.TIMESTAMP.equals(columnType)) {
-      columnNameTS = "TIMESTAMP(`" + columnNameTS + "`)";
-    }
-    if (partitionFromDate != null) {
-      timePartitionCondition.append(columnNameTS).append(" >= ").append("TIMESTAMP(\"")
-        .append(partitionFromDate).append("\")");
-    }
-    if (partitionFromDate != null && partitionToDate != null) {
-      timePartitionCondition.append(" and ");
-    }
-    if (partitionToDate != null) {
-      timePartitionCondition.append(columnNameTS).append(" < ").append("TIMESTAMP(\"")
-        .append(partitionToDate).append("\")");
-    }
-    return timePartitionCondition.toString();
-  }
 }