tg_mst.gsql

CREATE DISTRIBUTED QUERY tg_mst(VERTEX opt_source, SET<STRING> v_type_set, SET<STRING> e_type_set, STRING weight_attribute, STRING weight_type,
  INT maximum_iteration = -1, BOOL print_results = TRUE, STRING result_attribute = "", STRING file_path = "")  SYNTAX V1 {

    /*
    First Author: <First Author Name>
    First Commit Date:  <First Commit Date>

    Recent Author: <Recent Commit Author Name>
    Recent Commit Date: <Recent Commit Date>


    Repository:
        https://github.com/tigergraph/gsql-graph-algorithms/tree/master/algorithms/Path

    Maturity:
        production

    Description: 
        Returns a set of edges which form a Minumum Spanning Tree for a connected component. The algorithm
        cans tart either with a user-provided seed vertex or a randomly chosen one.  If you want a set of
        tree which span all the graph's components, use the msf (minimum spanning forest) algorithm.

    Publications:
        NA

    TigerGraph Documentation:
        https://docs.tigergraph.com/graph-ml/current/pathfinding-algorithms/minimum-spanning-tree

    Parameters:
        opt_source:
            start vertex (optional)
        print_results:
            If True, print JSON output
        v_type_set:
            vertex types to traverse
        result_attribute:
            INT attribute to store results to
        e_type_set:
            edge types to traverse
        file_path:
            file to write CSV output to
        weight_attribute:
            attribute for edge weights
        maximum_iteration:
            max iterations/edges (-1 = ALL)
        weight_type:
            weight data type (INT,FLOAT,DOUBLE)
    */

  TYPEDEF TUPLE<VERTEX from_v, VERTEX to_v, EDGE e, FLOAT weight, INT vid> EDGE_WEIGHT;
  HeapAccum<EDGE_WEIGHT>(1, weight ASC, vid ASC) @@chosen_edge_heap; // keep the minimal tuple
  SetAccum<EDGE_WEIGHT> @@mst_set;
  SetAccum<EDGE> @@result_set; 
  OrAccum @or_chosen;
  INT iter_limit;
  FILE f (file_path);
  
  # Check weight_type parameter
  IF weight_type NOT IN ("INT", "FLOAT", "DOUBLE") THEN
      PRINT "weight_type must be INT, FLOAT, or DOUBLE" AS errMsg;
      RETURN;
  END;


  # Pick the start vertex to initialize
  All_v = {v_type_set}; 
  MSTNodes = SELECT s 
             FROM All_v:s LIMIT 1;
  IF opt_source IS NOT NULL THEN
      MSTNodes = {opt_source};
  END;

  Current = SELECT s 
            FROM MSTNodes:s
	    POST-ACCUM s.@or_chosen = true;
  PRINT Current[Current.id] AS Source;

    
  # Find the MST 
  iter_limit = All_v.size();      # set max #iterations
  IF maximum_iteration > 0 THEN
      iter_limit = maximum_iteration;
  END;
  WHILE (Current.size() > 0) LIMIT iter_limit DO
      Current = SELECT t
	        FROM MSTNodes:s -(e_type_set:e)- v_type_set:t
	        WHERE t.@or_chosen == false    // vertex not in MSTNodes
	        ACCUM
	            CASE weight_type
		        WHEN "INT" THEN
		            @@chosen_edge_heap += EDGE_WEIGHT(s, t, e, e.getAttr(weight_attribute,"INT"), getvid(t))
		        WHEN "FLOAT" THEN
		            @@chosen_edge_heap += EDGE_WEIGHT(s, t, e, e.getAttr(weight_attribute,"FLOAT"), getvid(t))
		        WHEN "DOUBLE" THEN
			    @@chosen_edge_heap += EDGE_WEIGHT(s, t, e, e.getAttr(weight_attribute,"DOUBLE"), getvid(t))
		        END	
	        POST-ACCUM
	            IF t == @@chosen_edge_heap.top().to_v THEN    
		        t.@or_chosen = TRUE      // mark the chosen vertex to add into MSTNodes
		    END
	        HAVING t.@or_chosen == true;

      IF @@chosen_edge_heap.size() > 0 THEN
          IF result_attribute != "" THEN
	      S = SELECT s
	   	  FROM Current:s -(e_type_set:e) - v_type_set:t
		  WHERE t == @@chosen_edge_heap.top().from_v
		  ACCUM e.setAttr(result_attribute, TRUE);
	  END;
	  IF file_path != "" THEN
	      @@mst_set += @@chosen_edge_heap.top();
	  END;
	  IF print_results THEN
	      @@result_set += @@chosen_edge_heap.top().e;
	  END;
      END;    
      @@chosen_edge_heap.clear();
      MSTNodes = MSTNodes UNION Current;      // update MSTNodes
  END;

  # Output
  IF print_results THEN
      PRINT @@result_set as mst;
  END;

  IF file_path != "" THEN
      f.println("From", "To", "Weight");
      FOREACH e in @@mst_set DO
          f.println(e.from_v, e.to_v, e.weight);
      END;
  END;
}