-
Notifications
You must be signed in to change notification settings - Fork 966
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Write logged features to Offline Store (Go - Python integration) (
#2621) * write logged features by path Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * fix dummy sink Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * share snowflake cursor and stage name Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * share snowflake cursor and stage name Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * e2e tests WIP Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * graceful stop (grpc & logging) Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * revert accidental change Signed-off-by: pyalex <moskalenko.alexey@gmail.com> * add comments to go embedded methods Signed-off-by: Oleksii Moskalenko <oleksii@oleksiis-mbp.lan> Co-authored-by: Oleksii Moskalenko <oleksii@oleksiis-mbp.lan>
- Loading branch information
Showing
24 changed files
with
612 additions
and
114 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,83 @@ | ||
package logging | ||
|
||
import ( | ||
"errors" | ||
"fmt" | ||
"io" | ||
"io/ioutil" | ||
"os" | ||
"path/filepath" | ||
|
||
"github.com/apache/arrow/go/v8/arrow" | ||
"github.com/apache/arrow/go/v8/arrow/array" | ||
"github.com/apache/arrow/go/v8/parquet" | ||
"github.com/apache/arrow/go/v8/parquet/pqarrow" | ||
"github.com/google/uuid" | ||
) | ||
|
||
type OfflineStoreWriteCallback func(featureServiceName, datasetDir string) string | ||
|
||
type OfflineStoreSink struct { | ||
datasetDir string | ||
writeCallback OfflineStoreWriteCallback | ||
} | ||
|
||
func NewOfflineStoreSink(writeCallback OfflineStoreWriteCallback) (*OfflineStoreSink, error) { | ||
return &OfflineStoreSink{ | ||
datasetDir: "", | ||
writeCallback: writeCallback, | ||
}, nil | ||
} | ||
|
||
func (s *OfflineStoreSink) getOrCreateDatasetDir() (string, error) { | ||
if s.datasetDir != "" { | ||
return s.datasetDir, nil | ||
} | ||
dir, err := ioutil.TempDir("", "*") | ||
if err != nil { | ||
return "", err | ||
} | ||
s.datasetDir = dir | ||
return s.datasetDir, nil | ||
} | ||
|
||
func (s *OfflineStoreSink) cleanCurrentDatasetDir() error { | ||
if s.datasetDir == "" { | ||
return nil | ||
} | ||
datasetDir := s.datasetDir | ||
s.datasetDir = "" | ||
return os.RemoveAll(datasetDir) | ||
} | ||
|
||
func (s *OfflineStoreSink) Write(record arrow.Record) error { | ||
fileName, _ := uuid.NewUUID() | ||
datasetDir, err := s.getOrCreateDatasetDir() | ||
if err != nil { | ||
return err | ||
} | ||
|
||
var writer io.Writer | ||
writer, err = os.Create(filepath.Join(datasetDir, fmt.Sprintf("%s.parquet", fileName.String()))) | ||
if err != nil { | ||
return err | ||
} | ||
table := array.NewTableFromRecords(record.Schema(), []arrow.Record{record}) | ||
|
||
props := parquet.NewWriterProperties(parquet.WithDictionaryDefault(false)) | ||
arrProps := pqarrow.DefaultWriterProps() | ||
return pqarrow.WriteTable(table, writer, 1000, props, arrProps) | ||
} | ||
|
||
func (s *OfflineStoreSink) Flush(featureServiceName string) error { | ||
if s.datasetDir == "" { | ||
return nil | ||
} | ||
|
||
errMsg := s.writeCallback(featureServiceName, s.datasetDir) | ||
if errMsg != "" { | ||
return errors.New(errMsg) | ||
} | ||
|
||
return s.cleanCurrentDatasetDir() | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.