Skip to content

Commit 3d25b72

Browse files
author
Oleksii Moskalenko
authored
Create table in BigQuery if doesn't exists when new FeatureSetSpec arrived to IngestionJob (#815)
* create bq table if doesnt exists * dots * constants & unified job name
1 parent a46f7d1 commit 3d25b72

2 files changed

Lines changed: 57 additions & 16 deletions

File tree

‎core/src/main/java/feast/core/job/JobUpdateTask.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -238,7 +238,7 @@ String createJobId(Source source, String storeName) {
238238
String.format(
239239
"%s-%d-to-%s-%s",
240240
source.getTypeString(), Objects.hashCode(source.getConfig()), storeName, dateSuffix);
241-
return jobId.replaceAll("_store", "-");
241+
return jobId.replaceAll("_store", "-").toLowerCase();
242242
}
243243

244244
private void logAudit(Action action, Job job, String detail, Object... args) {

‎storage/connectors/bigquery/src/main/java/feast/storage/connectors/bigquery/writer/FeatureSetSpecToTableSchema.java‎

Lines changed: 56 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,14 @@
4242
import org.apache.beam.sdk.values.KV;
4343
import org.slf4j.Logger;
4444

45+
/**
46+
* Converts {@link feast.proto.core.FeatureSetProto.FeatureSetSpec} into BigQuery schema. Serializes
47+
* it into json-like format {@link TableSchema}. Fetches existing schema to merge existing fields
48+
* with new ones.
49+
*
50+
* <p>As a side effect this Operation may create bq table (if it doesn't exist) to make
51+
* bootstrapping faster
52+
*/
4553
public class FeatureSetSpecToTableSchema
4654
extends DoFn<KV<String, FeatureSetProto.FeatureSetSpec>, KV<String, TableSchema>> {
4755
private BigQuery bqService;
@@ -51,6 +59,12 @@ public class FeatureSetSpecToTableSchema
5159
private static final Logger log =
5260
org.slf4j.LoggerFactory.getLogger(FeatureSetSpecToTableSchema.class);
5361

62+
// Reserved columns
63+
public static final String EVENT_TIMESTAMP_COLUMN = "event_timestamp";
64+
public static final String CREATED_TIMESTAMP_COLUMN = "created_timestamp";
65+
public static final String INGESTION_ID_COLUMN = "ingestion_id";
66+
public static final String JOB_ID_COLUMN = "job_id";
67+
5468
// Column description for reserved fields
5569
public static final String BIGQUERY_EVENT_TIMESTAMP_FIELD_DESCRIPTION =
5670
"Event time for the FeatureRow";
@@ -76,19 +90,44 @@ public void processElement(
7690
@Element KV<String, FeatureSetProto.FeatureSetSpec> element,
7791
OutputReceiver<KV<String, TableSchema>> output,
7892
ProcessContext context) {
79-
Schema schema = createSchemaFromSpec(element.getValue(), element.getKey());
80-
output.output(KV.of(element.getKey(), serializeSchema(schema)));
93+
String specKey = element.getKey();
94+
95+
Table existingTable = getExistingTable(specKey);
96+
Schema schema = createSchemaFromSpec(element.getValue(), specKey, existingTable);
97+
98+
if (existingTable == null) {
99+
createTable(specKey, schema);
100+
}
101+
102+
output.output(KV.of(specKey, serializeSchema(schema)));
81103
}
82104

83-
private Table getExistingTable(String specKey) {
105+
private TableId generateTableId(String specKey) {
84106
TableDestination tableDestination = BigQuerySinkHelpers.getTableDestination(dataset, specKey);
85107
TableReference tableReference = BigQueryHelpers.parseTableSpec(tableDestination.getTableSpec());
86-
TableId tableId =
87-
TableId.of(
88-
tableReference.getProjectId(),
89-
tableReference.getDatasetId(),
90-
tableReference.getTableId());
91-
return bqService.getTable(tableId);
108+
return TableId.of(
109+
tableReference.getProjectId(), tableReference.getDatasetId(), tableReference.getTableId());
110+
}
111+
112+
private Table getExistingTable(String specKey) {
113+
return bqService.getTable(generateTableId(specKey));
114+
}
115+
116+
private void createTable(String specKey, Schema schema) {
117+
TimePartitioning timePartitioning =
118+
TimePartitioning.newBuilder(TimePartitioning.Type.DAY)
119+
.setField(EVENT_TIMESTAMP_COLUMN)
120+
.build();
121+
122+
StandardTableDefinition tableDefinition =
123+
StandardTableDefinition.newBuilder()
124+
.setTimePartitioning(timePartitioning)
125+
.setSchema(schema)
126+
.build();
127+
128+
TableInfo tableInfo = TableInfo.of(generateTableId(specKey), tableDefinition);
129+
130+
bqService.create(tableInfo);
92131
}
93132

94133
/**
@@ -98,12 +137,14 @@ private Table getExistingTable(String specKey) {
98137
*
99138
* @param spec FeatureSet spec that this table is for
100139
* @param specKey String for retrieving existing table
140+
* @param existingTable Table fetched from BQ. Fields from existing table used to merge with new
141+
* schema
101142
* @return {@link Schema} containing all tombstoned and active fields.
102143
*/
103-
private Schema createSchemaFromSpec(FeatureSetProto.FeatureSetSpec spec, String specKey) {
144+
private Schema createSchemaFromSpec(
145+
FeatureSetProto.FeatureSetSpec spec, String specKey, Table existingTable) {
104146
List<Field> fields = new ArrayList<>();
105147
log.info("Table {} will have the following fields:", specKey);
106-
Table existingTable = getExistingTable(specKey);
107148

108149
for (FeatureSetProto.EntitySpec entitySpec : spec.getEntitiesList()) {
109150
Field.Builder builder =
@@ -133,14 +174,14 @@ private Schema createSchemaFromSpec(FeatureSetProto.FeatureSetSpec spec, String
133174
Map<String, Pair<StandardSQLTypeName, String>>
134175
reservedFieldNameToPairOfStandardSQLTypeAndDescription =
135176
ImmutableMap.of(
136-
"event_timestamp",
177+
EVENT_TIMESTAMP_COLUMN,
137178
Pair.of(StandardSQLTypeName.TIMESTAMP, BIGQUERY_EVENT_TIMESTAMP_FIELD_DESCRIPTION),
138-
"created_timestamp",
179+
CREATED_TIMESTAMP_COLUMN,
139180
Pair.of(
140181
StandardSQLTypeName.TIMESTAMP, BIGQUERY_CREATED_TIMESTAMP_FIELD_DESCRIPTION),
141-
"ingestion_id",
182+
INGESTION_ID_COLUMN,
142183
Pair.of(StandardSQLTypeName.STRING, BIGQUERY_INGESTION_ID_FIELD_DESCRIPTION),
143-
"job_id",
184+
JOB_ID_COLUMN,
144185
Pair.of(StandardSQLTypeName.STRING, BIGQUERY_JOB_ID_FIELD_DESCRIPTION));
145186
for (Map.Entry<String, Pair<StandardSQLTypeName, String>> entry :
146187
reservedFieldNameToPairOfStandardSQLTypeAndDescription.entrySet()) {

0 commit comments

Comments
 (0)