Fix BigQuery query template to retrieve training data

core/src/main/java/feast/core/training/BigQueryDatasetTemplater.java

@@ -70,8 +70,10 @@ public String createQuery(
       throw new NoSuchElementException("features not found: " + featureIds);
     }
 
+    assert featureInfos.size() > 0;
     String tableId = getBqTableId(featureInfos.get(0));
-    Features features = new Features(featureIds, tableId);
+    String entityName = featureInfos.get(0).getEntity().getName();
+    Features features = new Features(featureIds, tableId, entityName);
 
     String startDateStr = formatDateString(startDate);
     String endDateStr = formatDateString(endDate);
@@ -117,9 +119,11 @@ static final class Features {
     final List<String> columns;
     final String tableId;
 
-    public Features(List<String> featureIds, String tableId) {
-      this.columns = featureIds.stream()
-          .map(f -> f.replace(".", "_"))
+    Features(List<String> featureIds, String tableId, String entityName) {
+      // columns represent the column name in BigQuery
+      // feature with id "myentity.myfeature" will be represented as column "myfeature" in BigQuery
+      columns = featureIds.stream()
+          .map(f -> f.replace(".", "_").replace(entityName + "_", ""))
           .collect(Collectors.toList());
       this.tableId = tableId;
     }

core/src/main/resources/templates/bq_training.tmpl

@@ -1,11 +1,9 @@
 SELECT
-  {{ feature_set.tableId }}.id,
-  {{ feature_set.tableId }}.event_timestamp
-  {% for feature in feature_set.columns -%}
-  ,{{ feature }}
-  {%- endfor %}
+  id,
+  event_timestamp{%- if feature_set.columns | length > 0 %},{%- endif %}
+  {{ feature_set.columns | join(',') }}
 FROM
-  {{ feature_set.tableId }}
+  `{{ feature_set.tableId }}`
 WHERE event_timestamp >= TIMESTAMP("{{ start_date }}") AND event_timestamp <= TIMESTAMP(DATETIME_ADD("{{ end_date }}", INTERVAL 1 DAY))
 {% if limit is not none -%}
 LIMIT {{ limit }}

core/src/test/java/feast/core/training/BigQueryDatasetTemplaterTest.java

@@ -96,6 +96,7 @@ public void shouldPassCorrectArgumentToTemplateEngine() {
         Timestamps.fromSeconds(Instant.parse("2019-01-01T00:00:00.00Z").getEpochSecond());
     int limit = 100;
     String featureId = "myentity.feature1";
+    String expectedColumnName = "feature1";
     String tableId = "project.dataset.myentity";
 
     when(featureInfoRespository.findAllById(any(List.class)))
@@ -123,7 +124,7 @@ public void shouldPassCorrectArgumentToTemplateEngine() {
 
     Features features = (Features) actualContext.get("feature_set");
     assertThat(features.getColumns().size(), equalTo(1));
-    assertThat(features.getColumns().get(0), equalTo(featureId.replace(".", "_")));
+    assertThat(features.getColumns().get(0), equalTo(expectedColumnName));
     assertThat(features.getTableId(), equalTo(tableId));
   }
 

core/src/test/resources/sql/expQuery1.sql

@@ -1,11 +1,11 @@
 SELECT
-    project.dataset.myentity.id,
-    project.dataset.myentity.event_timestamp ,
-    myentity_feature1,
-    myentity_feature2,
-    myentity_feature3
+    id,
+    event_timestamp,
+    feature1,
+    feature2,
+    feature3
 FROM
-    project.dataset.myentity
+    `project.dataset.myentity`
 WHERE
     event_timestamp >= TIMESTAMP("2018-01-02")
     AND event_timestamp <= TIMESTAMP(DATETIME_ADD("2018-01-30", INTERVAL 1 DAY)) LIMIT 100

core/src/test/resources/sql/expQuery2.sql

@@ -1,9 +1,9 @@
 SELECT
-    project.dataset.myentity.id,
-    project.dataset.myentity.event_timestamp ,
-    myentity_feature1
+    id,
+    event_timestamp,
+    feature1
 FROM
-    project.dataset.myentity
+    `project.dataset.myentity`
 WHERE
     event_timestamp >= TIMESTAMP("2018-01-02")
     AND event_timestamp <= TIMESTAMP(DATETIME_ADD("2018-01-30", INTERVAL 1 DAY)) LIMIT 1000