Integrating a Graph Database as a Machine Learning Feature Store Registry

Abstract

The rapid evolution of machine learning systems and their extended application in production environments have unveiled the need for robust and scalable feature management solutions. The feature store, a repository that centrally controls features used in model training and inference, is essential to these systems. Traditionally, feature stores have stored feature metadata and lineage data in SQL-based registries. This method cannot effectively model and query complex dependencies between features and entities, whereas there are alternative tools intended for such tasks. This study investigates whether a graph database can function as a feature store registry, leveraging its native capability to represent intricate relationships and lineage in a more natural manner. This work extends Feast, an open-source feature store that is widely used in industry, with a Neo4j-backed registry, building on previous research that emphasizes the strengths of graph databases in multi-hop traversals and relationship-based searches. Our approach models feature store objects as nodes, with their dependencies and lineage captured as relationships. New functionality was added to detect relationships between features, along with custom CLI commands designed to benefit from the descriptive power of graph queries. To evaluate the proposed solution, performance tests were conducted measuring execution time across graph-based, SQL-based and file-based registries. The results indicate that the graph-based registry excels in managing intricate, multi-hop relationships and is particularly effective for applications requiring deep dependency analysis. However, its usability and performance advantages are most pronounced when the query patterns involve complex relationship traversals, whereas for simpler, more direct queries its benefits may be less significant and may get overshadowed by the query planning overhead. In conclusion, this study demonstrates how graph databases may effectively represent the intricate relationships between information, improving the management and interpretability of machine learning systems. By offering empirical proof of the advantages and disadvantages of graph-based methods, it not only closes a gap in the feature store literature but also paves the way for further research into hybrid registry architectures that can dynamically balance the advantages of several backend systems and optimize graph operations.