query-orchestration/internal/document/text/childElements.go

package documenttext

import (
	"log/slog"

	"github.com/aws/aws-sdk-go-v2/service/textract/types"
	"github.com/google/uuid"
)

func (s *Service) GetDirectChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block) []uuid.UUID {
	lines := []uuid.UUID{}
	found := map[uuid.UUID]*bool{}

	for _, id := range s.getRelationshipIds(blockMap[id]) {
		if found[id] != nil && !*found[id] {
			continue
		}

		lines = append(lines, id)
		f := true
		found[id] = &f

		s.removeChildren(id, blockMap, &found)
	}

	output := []uuid.UUID{}
	for _, id := range lines {
		if found[id] == nil || !*found[id] {
			continue
		}

		output = append(output, id)
	}

	return output
}

func (s *Service) removeChildren(id uuid.UUID, blockMap map[uuid.UUID]types.Block, found *map[uuid.UUID]*bool) {
	for _, childId := range s.getRelationshipIds(blockMap[id]) {
		f := false
		(*found)[childId] = &f
		s.removeChildren(childId, blockMap, found)
	}
}

func (s *Service) GetBlockMap(blocksList []types.Block) map[uuid.UUID]types.Block {
	blocks := make(map[uuid.UUID]types.Block, len(blocksList))
	for _, block := range blocksList {
		if block.Id == nil {
			continue
		}
		id, err := uuid.Parse(*block.Id)
		if err != nil {
			slog.Warn("error parsing block id for block map", "error", err, "id", *block.Id)
			continue
		}

		blocks[id] = block
	}

	return blocks
}

func (s *Service) getRelationshipIds(block types.Block) []uuid.UUID {
	childIds := []uuid.UUID{}
	for _, relationship := range block.Relationships {
		for _, idStr := range relationship.Ids {
			id, err := uuid.Parse(idStr)
			if err != nil {
				slog.Error("invalid child uuid", "id", id, "error", err)
				continue
			}

			childIds = append(childIds, id)
		}
	}

	return childIds
}

func (s *Service) getRelationshipIdsByRel(block types.Block, relType types.RelationshipType) []uuid.UUID {
	childIds := []uuid.UUID{}
	for _, relationship := range block.Relationships {
		if relationship.Type != relType {
			continue
		}
		for _, idStr := range relationship.Ids {
			id, err := uuid.Parse(idStr)
			if err != nil {
				slog.Error("invalid child uuid", "id", id, "error", err)
				continue
			}

			childIds = append(childIds, id)
		}
	}

	return childIds
}