Repository: oom-ai/oomstore
Pull request #1096 (2022-01-26 09:27:52)

Jinghan/improve Join
This PR does:
- query entity rows table, get `min_unix_milli` and `max_unix_milli`
- use the time range to filter out unnecessary revision ranges in `Join`


Patch:
diff --git a/internal/database/offline/bigquery/join.go b/internal/database/offline/bigquery/join.go
index bd6debe3..b814e86c 100644
--- a/internal/database/offline/bigquery/join.go
+++ b/internal/database/offline/bigquery/join.go
@@ -24,11 +24,16 @@ func (db *DB) Join(ctx context.Context, opt offline.JoinOpt) (*types.JoinResult,
 	doJoinOpt := sqlutil.DoJoinOpt{
 		JoinOpt:             opt,
 		QueryResults:        bigqueryQueryResults,
+		QueryTableTimeRange: bigqueryQueryTimeRange,
 		ReadJoinResultQuery: READ_JOIN_RESULT_QUERY,
 	}
 	return sqlutil.DoJoin(ctx, dbOpt, doJoinOpt)
 }
 
+func bigqueryQueryTimeRange(ctx context.Context, dbOpt dbutil.DBOpt, tableName string) (*types.DataTableTimeRange, error) {
+	return getTableTimeRange(ctx, dbOpt, tableName)
+}
+
 func bigqueryQueryResults(ctx context.Context, dbOpt dbutil.DBOpt, query string, header dbutil.ColumnList, dropTableNames []string, backendType types.BackendType) (*types.JoinResult, error) {
 	rows, err := dbOpt.BigQueryDB.Query(query).Read(ctx)
 	if err != nil {
diff --git a/internal/database/offline/bigquery/store.go b/internal/database/offline/bigquery/store.go
index 515a3225..64c6feba 100644
--- a/internal/database/offline/bigquery/store.go
+++ b/internal/database/offline/bigquery/store.go
@@ -2,12 +2,9 @@ package bigquery
 
 import (
 	"context"
-	"fmt"
 
 	"cloud.google.com/go/bigquery"
 	"github.com/oom-ai/oomstore/pkg/errdefs"
-	"github.com/spf13/cast"
-	"google.golang.org/api/iterator"
 	"google.golang.org/api/option"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -44,54 +41,6 @@ func (db *DB) Ping(ctx context.Context) error {
 	return errdefs.WithStack(err)
 }
 
-func (db *DB) TableSchema(ctx context.Context, opt offline.TableSchemaOpt) (*types.DataTableSchema, error) {
-	q := fmt.Sprintf(`SELECT column_name, data_type FROM %s.INFORMATION_SCHEMA.COLUMNS WHERE table_name = "%s"`, db.datasetID, opt.TableName)
-	rows, err := db.Query(q).Read(ctx)
-	if err != nil {
-		return nil, errdefs.WithStack(err)
-	}
-	var schema types.DataTableSchema
-	for {
-		recordMap := make(map[string]bigquery.Value)
-		err = rows.Next(&recordMap)
-		if err == iterator.Done {
-			break
-		}
-		if err != nil {
-			return nil, errdefs.WithStack(err)
-		}
-		valueType, err := dbutil.ValueType(Backend, cast.ToString(recordMap["data_type"]))
-		if err != nil {
-			return nil, err
-		}
-		schema.Fields = append(schema.Fields, types.DataTableFieldSchema{
-			Name:      cast.ToString(recordMap["column_name"]),
-			ValueType: valueType,
-		})
-	}
-	if opt.CheckTimeRange {
-		q = fmt.Sprintf(`
-		SELECT
-			MIN(%s) AS %s,
-			MAX(%s) AS %s
-		FROM %s.%s`, "unix_milli", "min_unix_milli", "unix_milli", "max_unix_milli", db.datasetID, opt.TableName)
-		rows, err = db.Query(q).Read(ctx)
-		if err != nil {
-			return nil, errdefs.WithStack(err)
-		}
-		recordMap := make(map[string]bigquery.Value)
-		err = rows.Next(&recordMap)
-		if err != nil {
-			return nil, errdefs.WithStack(err)
-		}
-		schema.TimeRange = types.DataTableTimeRange{
-			MinUnixMilli: cast.ToInt64(recordMap["min_unix_milli"]),
-			MaxUnixMilli: cast.ToInt64(recordMap["max_unix_milli"]),
-		}
-	}
-	return &schema, nil
-}
-
 func (db *DB) Snapshot(ctx context.Context, opt offline.SnapshotOpt) error {
 	dbOpt := dbutil.DBOpt{
 		Backend:    Backend,
diff --git a/internal/database/offline/bigquery/table_schema.go b/internal/database/offline/bigquery/table_schema.go
new file mode 100644
index 00000000..940522f0
--- /dev/null
+++ b/internal/database/offline/bigquery/table_schema.go
@@ -0,0 +1,78 @@
+package bigquery
+
+import (
+	"context"
+	"fmt"
+
+	"cloud.google.com/go/bigquery"
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
+	"github.com/oom-ai/oomstore/internal/database/offline"
+	"github.com/oom-ai/oomstore/pkg/errdefs"
+	"github.com/oom-ai/oomstore/pkg/oomstore/types"
+	"github.com/spf13/cast"
+	"google.golang.org/api/iterator"
+)
+
+func (db *DB) TableSchema(ctx context.Context, opt offline.TableSchemaOpt) (*types.DataTableSchema, error) {
+	q := fmt.Sprintf(`SELECT column_name, data_type FROM %s.INFORMATION_SCHEMA.COLUMNS WHERE table_name = "%s"`, db.datasetID, opt.TableName)
+	rows, err := db.Query(q).Read(ctx)
+	if err != nil {
+		return nil, errdefs.WithStack(err)
+	}
+	var schema types.DataTableSchema
+	for {
+		recordMap := make(map[string]bigquery.Value)
+		err = rows.Next(&recordMap)
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			return nil, errdefs.WithStack(err)
+		}
+		valueType, err := dbutil.ValueType(Backend, cast.ToString(recordMap["data_type"]))
+		if err != nil {
+			return nil, err
+		}
+		schema.Fields = append(schema.Fields, types.DataTableFieldSchema{
+			Name:      cast.ToString(recordMap["column_name"]),
+			ValueType: valueType,
+		})
+	}
+	if opt.CheckTimeRange {
+		dbOpt := dbutil.DBOpt{
+			Backend:    Backend,
+			BigQueryDB: db.Client,
+			DatasetID:  &db.datasetID,
+		}
+		timeRange, err := getTableTimeRange(ctx, dbOpt, opt.TableName)
+		if err != nil {
+			return nil, err
+		}
+		schema.TimeRange = *timeRange
+	}
+	return &schema, nil
+}
+
+func getTableTimeRange(ctx context.Context, dbOpt dbutil.DBOpt, tableName string) (*types.DataTableTimeRange, error) {
+	q := fmt.Sprintf(`
+		SELECT
+			MIN(%s) AS %s,
+			MAX(%s) AS %s
+		FROM %s.%s`, "unix_milli", "min_unix_milli", "unix_milli", "max_unix_milli", *dbOpt.DatasetID, tableName)
+	rows, err := dbOpt.BigQueryDB.Query(q).Read(ctx)
+	if err != nil {
+		return nil, errdefs.WithStack(err)
+	}
+	recordMap := make(map[string]bigquery.Value)
+	err = rows.Next(&recordMap)
+	if err != nil {
+		return nil, errdefs.WithStack(err)
+	}
+	min := cast.ToInt64(recordMap["min_unix_milli"])
+	max := cast.ToInt64(recordMap["max_unix_milli"])
+
+	return &types.DataTableTimeRange{
+		MinUnixMilli: &min,
+		MaxUnixMilli: &max,
+	}, nil
+}
diff --git a/internal/database/offline/sqlutil/join.go b/internal/database/offline/sqlutil/join.go
index 8f3a70a3..bec13239 100644
--- a/internal/database/offline/sqlutil/join.go
+++ b/internal/database/offline/sqlutil/join.go
@@ -13,6 +13,7 @@ import (
 )
 
 type QueryResults func(ctx context.Context, dbOpt dbutil.DBOpt, query string, header dbutil.ColumnList, dropTableNames []string, backendType types.BackendType) (*types.JoinResult, error)
+type QueryTableTimeRange func(ctx context.Context, dbOpt dbutil.DBOpt, tableName string) (*types.DataTableTimeRange, error)
 
 func Join(ctx context.Context, db *sqlx.DB, opt offline.JoinOpt, backendType types.BackendType) (*types.JoinResult, error) {
 	dbOpt := dbutil.DBOpt{
@@ -22,6 +23,7 @@ func Join(ctx context.Context, db *sqlx.DB, opt offline.JoinOpt, backendType typ
 	doJoinOpt := DoJoinOpt{
 		JoinOpt:             opt,
 		QueryResults:        sqlxQueryResults,
+		QueryTableTimeRange: sqlxQueryTableTimeRange,
 		ReadJoinResultQuery: READ_JOIN_RESULT_QUERY,
 	}
 	return DoJoin(ctx, dbOpt, doJoinOpt)
@@ -30,6 +32,7 @@ func Join(ctx context.Context, db *sqlx.DB, opt offline.JoinOpt, backendType typ
 type DoJoinOpt struct {
 	offline.JoinOpt
 	QueryResults        QueryResults
+	QueryTableTimeRange QueryTableTimeRange
 	ReadJoinResultQuery string
 }
 
@@ -55,6 +58,15 @@ func DoJoin(ctx context.Context, dbOpt dbutil.DBOpt, opt DoJoinOpt) (*types.Join
 	if err != nil {
 		return nil, err
 	}
+	timeRange, err := opt.QueryTableTimeRange(ctx, dbOpt, entityRowsTableName)
+	if err != nil {
+		return nil, err
+	}
+	if timeRange.MinUnixMilli == nil || timeRange.MaxUnixMilli == nil {
+		data := make(chan []interface{})
+		defer close(data)
+		return &types.JoinResult{Data: data}, nil
+	}
 
 	// Step 2: process features by group, insert result to table joined
 	tableNames := make([]string, 0)
@@ -79,6 +91,7 @@ func DoJoin(ctx context.Context, dbOpt dbutil.DBOpt, opt DoJoinOpt) (*types.Join
 			Features:            featureList,
 			RevisionRanges:      revisionRanges,
 			EntityRowsTableName: entityRowsTableName,
+			TimeRange:           *timeRange,
 		})
 		if err != nil {
 			return nil, err
@@ -112,6 +125,7 @@ type joinOneGroupOpt struct {
 	EntityName          string
 	EntityRowsTableName string
 	ValueNames          []string
+	TimeRange           types.DataTableTimeRange
 }
 
 func joinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt joinOneGroupOpt) ([]string, error) {
@@ -128,6 +142,9 @@ func joinOneGroup(ctx context.Context, dbOpt dbutil.DBOpt, opt joinOneGroupOpt)
 	// Step 2: iterate each table range, join entity_rows table and each data tables
 	columns := append(opt.ValueNames, opt.Features.Names()...)
 	for _, r := range opt.RevisionRanges {
+		if *opt.TimeRange.MaxUnixMilli < r.MinRevision || *opt.TimeRange.MinUnixMilli > r.MaxRevision {
+			continue
+		}
 		query, err := buildJoinQuery(joinQueryParams{
 			TableName:           snapshotJoinedTableName,
 			EntityName:          opt.EntityName,
@@ -322,3 +339,7 @@ func sqlxQueryResults(ctx context.Context, dbOpt dbutil.DBOpt, query string, hea
 		Data:   data,
 	}, dropErr
 }
+
+func sqlxQueryTableTimeRange(ctx context.Context, dbOpt dbutil.DBOpt, tableName string) (*types.DataTableTimeRange, error) {
+	return getCdcTimeRange(ctx, dbOpt.SqlxDB, tableName, dbOpt.Backend)
+}
diff --git a/pkg/oomstore/import_stream.go b/pkg/oomstore/import_stream.go
index b6335459..06231dd6 100644
--- a/pkg/oomstore/import_stream.go
+++ b/pkg/oomstore/import_stream.go
@@ -114,7 +114,7 @@ func (s *OomStore) tableLinkImportStream(ctx context.Context, opt types.ImportSt
 	}
 
 	_, _, err = s.createRevision(ctx, metadata.CreateRevisionOpt{
-		Revision:    tableSchema.TimeRange.MinUnixMilli,
+		Revision:    *tableSchema.TimeRange.MinUnixMilli,
 		GroupID:     group.ID,
 		CdcTable:    &dataSource.TableName,
 		Description: opt.Description,
@@ -194,9 +194,11 @@ func (s *OomStore) validateRevisions(ctx context.Context, groupID int, schema *t
 	if err != nil {
 		return err
 	}
-
-	revisionBeforeMin := revisions.Before(schema.TimeRange.MinUnixMilli)
-	revisionBeforeMax := revisions.Before(schema.TimeRange.MaxUnixMilli)
+	if schema.TimeRange.MinUnixMilli == nil || schema.TimeRange.MaxUnixMilli == nil {
+		return errdefs.Errorf("table is empty")
+	}
+	revisionBeforeMin := revisions.Before(*schema.TimeRange.MinUnixMilli)
+	revisionBeforeMax := revisions.Before(*schema.TimeRange.MaxUnixMilli)
 	if revisionBeforeMax != revisionBeforeMin {
 		return errdefs.Errorf("data table crosses with another offline table %s", revisionBeforeMax.CdcTable)
 	}
@@ -208,7 +210,7 @@ func (s *OomStore) validateRevisions(ctx context.Context, groupID int, schema *t
 		if err != nil {
 			return err
 		}
-		if beforeTableSchema.TimeRange.MaxUnixMilli >= schema.TimeRange.MinUnixMilli {
+		if beforeTableSchema.TimeRange.MaxUnixMilli != nil && *beforeTableSchema.TimeRange.MaxUnixMilli >= *schema.TimeRange.MinUnixMilli {
 			return errdefs.Errorf("data table crosses with another offline table %s", revisionBeforeMin.CdcTable)
 		}
 	}
diff --git a/pkg/oomstore/types/types.go b/pkg/oomstore/types/types.go
index 1b391e71..b0b1902c 100644
--- a/pkg/oomstore/types/types.go
+++ b/pkg/oomstore/types/types.go
@@ -67,8 +67,8 @@ func (e *ExportResult) GetErrorChannel() <-chan error {
 }
 
 type DataTableTimeRange struct {
-	MinUnixMilli int64 `db:"min_unix_milli"`
-	MaxUnixMilli int64 `db:"max_unix_milli"`
+	MinUnixMilli *int64 `db:"min_unix_milli"`
+	MaxUnixMilli *int64 `db:"max_unix_milli"`
 }
 
 type DataTableSchema struct {


Test patch:
diff --git a/internal/database/offline/snowflake/store_test.go b/internal/database/offline/snowflake/store_test.go
index 1f0485d3..0855be67 100644
--- a/internal/database/offline/snowflake/store_test.go
+++ b/internal/database/offline/snowflake/store_test.go
@@ -7,6 +7,8 @@ import (
 	"strings"
 	"testing"
 
+	"github.com/oom-ai/oomstore/internal/database/offline/test_impl"
+
 	"github.com/stretchr/testify/require"
 
 	"github.com/oom-ai/oomstore/internal/database/dbutil"
@@ -53,25 +55,25 @@ func destroyStore(database string) func() {
 	}
 }
 
-//func TestPing(t *testing.T) {
-//	test_impl.TestPing(t, prepareStore, destroyStore(DATABASE))
-//}
-//
-//func TestExport(t *testing.T) {
-//	test_impl.TestExport(t, prepareStore, destroyStore(DATABASE))
-//}
-//
-//func TestImport(t *testing.T) {
-//	test_impl.TestImport(t, prepareStore, destroyStore(DATABASE))
-//}
-//
-//func TestJoin(t *testing.T) {
-//	test_impl.TestJoin(t, prepareStore, destroyStore(DATABASE))
-//}
-//
-//func TestSnapshot(t *testing.T) {
-//	test_impl.TestSnapshot(t, prepareStore, destroyStore(DATABASE))
-//}
+func TestPing(t *testing.T) {
+	test_impl.TestPing(t, prepareStore, destroyStore(DATABASE))
+}
+
+func TestExport(t *testing.T) {
+	test_impl.TestExport(t, prepareStore, destroyStore(DATABASE))
+}
+
+func TestImport(t *testing.T) {
+	test_impl.TestImport(t, prepareStore, destroyStore(DATABASE))
+}
+
+func TestJoin(t *testing.T) {
+	test_impl.TestJoin(t, prepareStore, destroyStore(DATABASE))
+}
+
+func TestSnapshot(t *testing.T) {
+	test_impl.TestSnapshot(t, prepareStore, destroyStore(DATABASE))
+}
 
 // We don't use test_impl.TestTableSchema because snowflake cannot be
 // accessed by two different sessions.
@@ -112,8 +114,8 @@ func TestTableSchema(t *testing.T) {
 			},
 		},
 		TimeRange: types.DataTableTimeRange{
-			MinUnixMilli: 1,
-			MaxUnixMilli: 100,
+			MinUnixMilli: int64Ptr(1),
+			MaxUnixMilli: int64Ptr(100),
 		},
 	}
 	require.ElementsMatch(t, expected.Fields, actual.Fields)
@@ -121,6 +123,10 @@ func TestTableSchema(t *testing.T) {
 
 }
 
-//func TestCreateTable(t *testing.T) {
-//	test_impl.TestCreateTable(t, prepareStore, destroyStore(DATABASE))
-//}
+func int64Ptr(i int64) *int64 {
+	return &i
+}
+
+func TestCreateTable(t *testing.T) {
+	test_impl.TestCreateTable(t, prepareStore, destroyStore(DATABASE))
+}
diff --git a/internal/database/offline/test_impl/table_schema.go b/internal/database/offline/test_impl/table_schema.go
index e76c38f4..5b44ab7e 100644
--- a/internal/database/offline/test_impl/table_schema.go
+++ b/internal/database/offline/test_impl/table_schema.go
@@ -4,6 +4,8 @@ import (
 	"context"
 	"testing"
 
+	"github.com/stretchr/testify/assert"
+
 	"github.com/oom-ai/oomstore/internal/database/offline"
 
 	"github.com/oom-ai/oomstore/pkg/oomstore/types"
@@ -41,9 +43,14 @@ func TestTableSchema(t *testing.T, prepareStore PrepareStoreFn, destroyStore Des
 			},
 		},
 		TimeRange: types.DataTableTimeRange{
-			MinUnixMilli: 10,
-			MaxUnixMilli: 100,
+			MinUnixMilli: int64Ptr(1),
+			MaxUnixMilli: int64Ptr(100),
 		},
 	}
-	require.ElementsMatch(t, expected.Fields, actual.Fields)
+	assert.ElementsMatch(t, expected.Fields, actual.Fields)
+	assert.Equal(t, expected.TimeRange, actual.TimeRange)
+}
+
+func int64Ptr(i int64) *int64 {
+	return &i
 }