Use trips as primary unit for filtering, cascade changes upwards, and remove nonexistent association

lib/gtfs_df/feed.rb

@@ -36,7 +36,7 @@ class Feed
       booking_rules
     ].freeze
 
-    attr_reader(*GTFS_FILES)
+    attr_reader(*GTFS_FILES, :graph)
 
     # Initialize with a hash of DataFrames
     REQUIRED_GTFS_FILES = %w[agency stops routes trips stop_times].freeze
@@ -53,6 +53,8 @@ def initialize(data = {})
         end.join(", ")}"
       end
 
+      @graph = GtfsDf::Graph.build
+
       GTFS_FILES.each do |file|
         df = data[file]
         schema_class_name = file.split("_").map(&:capitalize).join
@@ -68,81 +70,36 @@ def initialize(data = {})
       end
     end
 
-    # Load from a directory of GTFS CSV files
-    def self.load_from_dir(dir)
-      data = {}
-      GTFS_FILES.each do |file|
-        path = File.join(dir, "#{file}.txt")
-        next unless File.exist?(path)
-
-        schema_class_name = file.split("_").map(&:capitalize).join
-
-        data[file] = GtfsDf::Schema.const_get(schema_class_name).new(path)
-      end
-      new(data)
-    end
-
     # Filter the feed using a view hash
     # Example view: { 'routes' => { 'route_id' => '123' }, 'trips' => { 'service_id' => 'A' } }
     def filter(view)
       filtered = {}
-      graph = GtfsDf::Graph.build
-      # Step 1: Apply view filters
+
       GTFS_FILES.each do |file|
         df = send(file)
         next unless df
 
-        filters = view[file]
-        if filters && !filters.empty?
-          filters.each do |col, val|
-            df = if val.is_a?(Array)
-              df.filter(Polars.col(col).is_in(val))
-            elsif val.respond_to?(:call)
-              df.filter(val.call(Polars.col(col)))
-            else
-              df.filter(Polars.col(col).eq(val))
-            end
-          end
-        end
         filtered[file] = df
       end
-      # Step 2: Cascade filters following the directed edges
-      # An edge from parent->child means: filter child based on valid parent IDs
-      changed = true
-      while changed
-        changed = false
-        GTFS_FILES.each do |parent_file|
-          parent_df = filtered[parent_file]
-          next unless parent_df
-
-          # For each outgoing edge from parent_file to child_file
-          graph.adj[parent_file]&.each do |child_file, attrs|
-            child_df = filtered[child_file]
-            next unless child_df && child_df.height > 0
-
-            attrs[:dependencies].each do |dep|
-              parent_col = dep[parent_file]
-              child_col = dep[child_file]
-
-              next unless parent_col && child_col &&
-                parent_df.columns.include?(parent_col) && child_df.columns.include?(child_col)
-
-              # Get valid values from parent
-              valid_values = parent_df[parent_col].to_a.uniq.compact
-
-              # Filter child to only include rows that reference valid parent values
-              before = child_df.height
-              child_df = child_df.filter(Polars.col(child_col).is_in(valid_values))
-
-              if child_df.height < before
-                filtered[child_file] = child_df
-                changed = true
-              end
-            end
-          end
+
+      # Trips are the atomic unit of GTFS, we will generate a new view
+      # based on the set of trips that would be included for each invidual filter
+      # and cascade changes from this view in order to retain referential integrity
+      trip_ids = nil
+
+      view.each do |file, filters|
+        new_filtered = filter!(file, filters, filtered.dup)
+        trip_ids = if trip_ids.nil?
+          new_filtered["trips"]["trip_id"]
+        else
+          trip_ids & new_filtered["trips"]["trip_id"]
         end
       end
 
+      if trip_ids
+        filtered = filter!("trips", {"trip_id" => trip_ids.to_a}, filtered)
+      end
+
       # Remove files that are empty, but keep required files even if empty
       filtered.delete_if do |file, df|
         is_required_file = REQUIRED_GTFS_FILES.include?(file) ||
@@ -153,5 +110,60 @@ def filter(view)
       end
       self.class.new(filtered)
     end
+
+    private
+
+    def filter!(file, filters, filtered)
+      unless filters.empty?
+        df = filtered[file]
+
+        filters.each do |col, val|
+          df = if val.is_a?(Array)
+            df.filter(Polars.col(col).is_in(val))
+          elsif val.respond_to?(:call)
+            df.filter(val.call(Polars.col(col)))
+          else
+            df.filter(Polars.col(col).eq(val))
+          end
+        end
+
+        filtered[file] = df
+
+        prune!(file, filtered)
+      end
+
+      filtered
+    end
+
+    def prune!(root, filtered)
+      graph.each_bfs_edge(root) do |parent_file, child_file|
+        parent_df = filtered[parent_file]
+        next unless parent_df
+
+        child_df = filtered[child_file]
+        next unless child_df && child_df.height > 0
+
+        attrs = graph.get_edge_data(parent_file, child_file)
+
+        attrs[:dependencies].each do |dep|
+          parent_col = dep[parent_file]
+          child_col = dep[child_file]
+
+          next unless parent_col && child_col &&
+            parent_df.columns.include?(parent_col) && child_df.columns.include?(child_col)
+
+          # Get valid values from parent
+          valid_values = parent_df[parent_col].to_a.uniq.compact
+
+          # Filter child to only include rows that reference valid parent values
+          before = child_df.height
+          child_df = child_df.filter(Polars.col(child_col).is_in(valid_values))
+
+          if child_df.height < before
+            filtered[child_file] = child_df
+          end
+        end
+      end
+    end
   end
 end

lib/gtfs_df/graph.rb

@@ -4,15 +4,15 @@ module GtfsDf
   class Graph
     # Returns a directed graph of GTFS file dependencies
     def self.build
-      g = NetworkX::DiGraph.new
+      g = NetworkX::Graph.new
       # Nodes: GTFS files
       files = %w[
         agency routes trips stop_times stops calendar calendar_dates shapes transfers frequencies fare_attributes fare_rules
         fare_leg_join_rules fare_transfer_rules areas networks route_networks location_groups location_group_stops booking_rules
       ]
       files.each { |f| g.add_node(f) }
 
-      # Edges: dependencies
+      # TODO: Add fare_rules -> stops + test
       edges = [
         ["agency", "routes", {dependencies: [
           {"agency" => "agency_id", "routes" => "agency_id"}
@@ -116,9 +116,6 @@ def self.build
         ["booking_rules", "stop_times", {dependencies: [
           {"booking_rules" => "booking_rule_id", "stop_times" => "pickup_booking_rule_id"},
           {"booking_rules" => "booking_rule_id", "stop_times" => "drop_off_booking_rule_id"}
-        ]}],
-        ["stops", "booking_rules", {dependencies: [
-          {"stops" => "stop_id", "booking_rules" => "stop_id"}
         ]}]
       ]