Move StructuredSubgraph alongside batch merge function.

`graphstore` package owns entire batch merge operation.
Added tests for node and rel batching.
This commit is contained in:
Jay
2026-03-03 17:24:39 -05:00
parent 4fddf849d0
commit e43ed4af55
4 changed files with 343 additions and 297 deletions

View File

@@ -6,7 +6,6 @@ package graph
import (
"fmt"
"sort"
"strings"
)
// ========================================
@@ -160,167 +159,3 @@ func (r *Relationship) Serialize() *SerializedRel {
srel["end"] = r.End.Props
return &srel
}
// ========================================
// Structured Subgraph
// ========================================
// StructuredSubgraph is a structured collection of nodes and relationships for
// the purpose of conducting batch operations.
type StructuredSubgraph struct {
// A map of grouped nodes, batched by their label combinations.
nodes map[string][]*Node
// A map of grouped relationships, batched by their type and related node
// labels.
rels map[string][]*Relationship
// Provides node property keys used to match nodes with given labels in the
// database.
matchProvider MatchKeysProvider
}
// NewStructuredSubgraph creates an empty structured subgraph with the given
// match keys provider.
func NewStructuredSubgraph(matchProvider MatchKeysProvider) *StructuredSubgraph {
return &StructuredSubgraph{
nodes: make(map[string][]*Node),
rels: make(map[string][]*Relationship),
matchProvider: matchProvider,
}
}
// AddNode adds a node into the subgraph.
func (s *StructuredSubgraph) AddNode(node *Node) error {
// Verify that the node has defined match property values.
matchLabel, _, err := node.MatchProps(s.matchProvider)
if err != nil {
return fmt.Errorf("invalid node: %s", err)
}
// Determine the node's batch key.
batchKey := createNodeBatchKey(matchLabel, node.Labels.ToArray())
if _, exists := s.nodes[batchKey]; !exists {
s.nodes[batchKey] = []*Node{}
}
// Add the node to the subgraph.
s.nodes[batchKey] = append(s.nodes[batchKey], node)
return nil
}
// AddRel adds a relationship into the subgraph.
func (s *StructuredSubgraph) AddRel(rel *Relationship) error {
// Verify that the start node has defined match property values.
startLabel, _, err := rel.Start.MatchProps(s.matchProvider)
if err != nil {
return fmt.Errorf("invalid start node: %s", err)
}
// Verify that the end node has defined match property values.
endLabel, _, err := rel.End.MatchProps(s.matchProvider)
if err != nil {
return fmt.Errorf("invalid end node: %s", err)
}
// Determine the relationship's batch key.
batchKey := createRelBatchKey(rel.Type, startLabel, endLabel)
if _, exists := s.rels[batchKey]; !exists {
s.rels[batchKey] = []*Relationship{}
}
// Add the relationship to the subgraph.
s.rels[batchKey] = append(s.rels[batchKey], rel)
return nil
}
// GetNodes returns the nodes grouped under the given batch key.
func (s *StructuredSubgraph) GetNodes(nodeKey string) []*Node {
return s.nodes[nodeKey]
}
// GetRels returns the rels grouped under the given batch key.
func (s *StructuredSubgraph) GetRels(relKey string) []*Relationship {
return s.rels[relKey]
}
func (s *StructuredSubgraph) MatchProvider() MatchKeysProvider {
return s.matchProvider
}
// NodeCount returns the number of nodes in the subgraph.
func (s *StructuredSubgraph) NodeCount() int {
count := 0
for l := range s.nodes {
count += len(s.nodes[l])
}
return count
}
// RelCount returns the number of relationships in the subgraph.
func (s *StructuredSubgraph) RelCount() int {
count := 0
for t := range s.rels {
count += len(s.rels[t])
}
return count
}
// NodeKeys returns the list of node batch keys in the subgraph.
func (s *StructuredSubgraph) NodeKeys() []string {
keys := []string{}
for l := range s.nodes {
keys = append(keys, l)
}
return keys
}
// RelKeys returns the list of relationship batch keys in the subgraph.
func (s *StructuredSubgraph) RelKeys() []string {
keys := []string{}
for t := range s.rels {
keys = append(keys, t)
}
return keys
}
// createNodeBatchKey returns the serialized node labels for batching.
func createNodeBatchKey(matchLabel string, labels []string) string {
sort.Strings(labels)
serializedLabels := strings.Join(labels, ",")
return fmt.Sprintf("%s:%s", matchLabel, serializedLabels)
}
// createRelBatchKey returns the serialized relationship type and start/end node
// labels for batching.
func createRelBatchKey(
rtype string, startLabel string, endLabel string) string {
return strings.Join([]string{rtype, startLabel, endLabel}, ",")
}
// DeserializeNodeBatchKey returns the list of node labels from the serialized batch
// key.
func DeserializeNodeBatchKey(batchKey string) (string, []string, error) {
parts := strings.Split(batchKey, ":")
if len(parts) != 2 {
return "", nil, fmt.Errorf("invalid node batch key: %s", batchKey)
}
matchLabel, serializedLabels := parts[0], parts[1]
labels := strings.Split(serializedLabels, ",")
return matchLabel, labels, nil
}
// DeserializeRelBatchKey returns the relationship type, start node label, and end
// node label from the serialized batch key. Panics if the batch key is invalid.
func DeserializeRelBatchKey(batchKey string) (string, string, string, error) {
parts := strings.Split(batchKey, ",")
if len(parts) != 3 {
return "", "", "", fmt.Errorf("invalid relationship batch key: %s", batchKey)
}
rtype, startLabel, endLabel := parts[0], parts[1], parts[2]
return rtype, startLabel, endLabel, nil
}