swiss-netex 1.0.0

data/lib/swiss_netex/filter/timetable.rb

@@ -0,0 +1,387 @@
+# frozen_string_literal: true
+
+require "nokogiri"
+require "stringio"
+require "tempfile"
+
+module SwissNetex
+  class Filter
+    # Stream one TIMETABLE shard in a single pass:
+    # - keep matching ServiceJourneys (collect refs)
+    # - keep TrainNumber definitions referenced by those journeys
+    #
+    # Dropped journeys never materialize their (large) calls list. Kept journeys
+    # write the head buffer then stream the body so peak strings stay short.
+    # Empty shards return io: nil so the caller omits them from the output zip.
+    #
+    # TrainNumber defs may appear before or after vehicleJourneys. When the
+    # national export puts them after journeys, matching defs write in place.
+    # When defs come first (or no journey refs are known yet), the block is
+    # buffered and flushed after </vehicleJourneys> (or TimetableFrame end).
+    # That can reorder trainNumbers after journeys — intentional and small.
+    #
+    # Pure whitespace before dropped journeys/trains is discarded so output
+    # does not keep blank holes between kept siblings (same idea as XmlSieve).
+    class Timetable
+      JOURNEY_NAMES = %w[ServiceJourney TemplateServiceJourney].freeze
+      JOURNEY_SET = JOURNEY_NAMES.to_h { |name| [name, true] }.freeze
+
+      TYPE_ELEMENT = Nokogiri::XML::Reader::TYPE_ELEMENT
+      TYPE_END = Nokogiri::XML::Reader::TYPE_END_ELEMENT
+      TYPE_TEXT = Nokogiri::XML::Reader::TYPE_TEXT
+      TYPE_CDATA = Nokogiri::XML::Reader::TYPE_CDATA
+      TYPE_WS = Nokogiri::XML::Reader::TYPE_WHITESPACE
+      TYPE_SWS = Nokogiri::XML::Reader::TYPE_SIGNIFICANT_WHITESPACE
+      TYPE_COMMENT = Nokogiri::XML::Reader::TYPE_COMMENT
+
+      Result = Data.define(:io, :journeys_kept, :train_numbers_kept)
+
+      # line_filter: true when the user passed explicit --line aliases. In that
+      # mode journeys without a matching LineRef are dropped (no operator-only
+      # fallback). When false, selected operator journeys without LineRef stay.
+      def initialize(operator_ids:, line_ids:, refs:, line_filter: false)
+        @operator_ids = operator_ids.to_h { |id| [id.to_s, true] }
+        @line_ids = line_ids.to_h { |id| [id.to_s, true] }
+        @line_filter = line_filter
+        @refs = refs
+        @deferred_trains = []
+      end
+
+      def call(source)
+        out = new_tempfile
+        out << "<?xml version=\"1.0\" encoding=\"utf-8\"?>\n"
+        journeys_kept = 0
+        train_numbers_kept = 0
+        @deferred_trains = []
+        @journeys_seen = false
+        @pending_ws = +""
+        reader = Nokogiri::XML::Reader(to_io(source))
+
+        while reader.read
+          case reader.node_type
+          when TYPE_ELEMENT
+            start = handle_start(reader, out)
+            case start
+            when :kept_journey
+              journeys_kept += 1
+            when Array
+              train_numbers_kept += start[1] if start[0] == :kept_trains
+            end
+          when TYPE_END
+            name = reader.local_name
+            if name == "TimetableFrame"
+              train_numbers_kept += flush_deferred_trains(out)
+              flush_pending_ws(out)
+              out << "</#{reader.name}>"
+            else
+              flush_pending_ws(out)
+              out << "</#{reader.name}>"
+              train_numbers_kept += flush_deferred_trains(out) if name == "vehicleJourneys"
+            end
+          when TYPE_WS, TYPE_SWS
+            buffer_ws(reader.value.to_s)
+          when TYPE_TEXT
+            handle_text(reader, out)
+          when TYPE_CDATA
+            flush_pending_ws(out)
+            out << reader.value.to_s
+          when TYPE_COMMENT
+            flush_pending_ws(out)
+            out << "<!--#{reader.value}-->"
+          end
+        end
+
+        train_numbers_kept += flush_deferred_trains(out)
+        flush_pending_ws(out)
+
+        if journeys_kept.zero?
+          close_io(out)
+          return Result.new(io: nil, journeys_kept: 0, train_numbers_kept: 0)
+        end
+
+        out.flush
+        out.rewind
+        Result.new(io: out, journeys_kept: journeys_kept, train_numbers_kept: train_numbers_kept)
+      end
+
+      private
+
+      def handle_start(reader, out)
+        name = reader.local_name
+        if JOURNEY_SET[name]
+          result = write_or_skip_journey(reader, out)
+          @journeys_seen = true if %i[kept_journey dropped].include?(result)
+          clear_pending_ws if result == :dropped
+          result
+        elsif name == "trainNumbers"
+          result = collect_train_numbers(reader, out)
+          clear_pending_ws if %i[dropped deferred].include?(result)
+          result
+        elsif name == "TrainNumber"
+          result = buffer_or_write_train(reader, out)
+          clear_pending_ws if %i[dropped deferred].include?(result)
+          result
+        else
+          flush_pending_ws(out)
+          write_start_to(reader, out)
+          :structure
+        end
+      end
+
+      # Walk direct children until LineRef/OperatorRef decide keep/drop.
+      # On keep: flush head, stream remaining children to out while indexing refs.
+      def write_or_skip_journey(reader, out)
+        depth = reader.depth
+        head = +""
+        journey_id = reader.attribute("id")
+        resp = reader.attribute("responsibilitySetRef")
+        write_start_to(reader, head)
+        return :dropped if reader.empty_element?
+
+        line_ref = nil
+        op_ref = nil
+        keep = nil
+
+        while reader.read
+          if reader.node_type == TYPE_END && reader.depth == depth
+            head << "</#{reader.name}>"
+            break
+          end
+
+          if reader.node_type == TYPE_ELEMENT && reader.depth == depth + 1
+            child_name = reader.local_name
+
+            if child_name == "LineRef"
+              line_ref = reader.attribute("ref")
+              keep = line_ref && @line_ids.key?(line_ref)
+              unless keep
+                skip_to_depth(reader, depth)
+                return :dropped
+              end
+
+              child_xml = reader.outer_xml
+              skip_element(reader)
+              head << child_xml
+              flush_kept_journey(out, head, journey_id, resp)
+              stream_journey_body(reader, depth, out)
+              return :kept_journey
+            end
+
+            op_ref = reader.attribute("ref") if child_name == "OperatorRef"
+
+            child_xml = reader.outer_xml
+            skip_element(reader)
+            head << child_xml
+          else
+            append_leaf(reader, head)
+          end
+        end
+
+        keep = operator_only_keep?(op_ref) if keep.nil?
+        return :dropped unless keep
+
+        flush_kept_journey(out, head, journey_id, resp)
+        :kept_journey
+      end
+
+      def flush_kept_journey(out, head, journey_id, resp)
+        flush_pending_ws(out)
+        out << head
+        @refs.add(:journeys, journey_id) if journey_id
+        @refs.add(:responsibility_sets, resp) if resp
+        @refs.merge_from_xml(head)
+      end
+
+      def stream_journey_body(reader, depth, out)
+        while reader.read
+          if reader.node_type == TYPE_END && reader.depth == depth
+            out << "</#{reader.name}>"
+            break
+          end
+
+          if reader.node_type == TYPE_ELEMENT && reader.depth == depth + 1
+            child_xml = reader.outer_xml
+            skip_element(reader)
+            out << child_xml
+            @refs.merge_from_xml(child_xml)
+          else
+            append_leaf(reader, out)
+          end
+        end
+      end
+
+      # Buffer every TrainNumber in the container. Write filtered container now
+      # when journey refs already exist; otherwise defer until after journeys.
+      def collect_train_numbers(reader, out)
+        trains = read_train_number_children(reader)
+        write_or_defer_trains(trains, out)
+      end
+
+      def buffer_or_write_train(reader, out)
+        id = reader.attribute("id")
+        xml = reader.outer_xml
+        skip_element(reader)
+        return :dropped unless id
+
+        write_or_defer_trains([[id, xml]], out)
+      end
+
+      def write_or_defer_trains(trains, out)
+        return :dropped if trains.empty?
+
+        kept = matching_trains(trains)
+        if kept.any?
+          flush_pending_ws(out)
+          write_train_numbers(out, kept)
+          return [:kept_trains, kept.size]
+        end
+
+        # After journeys, unmatched defs are dead; only buffer when we may still
+        # learn TrainNumberRefs from later vehicleJourneys.
+        return :dropped if @journeys_seen
+
+        @deferred_trains.concat(trains)
+        :deferred
+      end
+
+      def flush_deferred_trains(out)
+        return 0 if @deferred_trains.empty?
+
+        trains = @deferred_trains
+        @deferred_trains = []
+        kept = matching_trains(trains)
+        return 0 if kept.empty?
+
+        flush_pending_ws(out)
+        write_train_numbers(out, kept)
+        kept.size
+      end
+
+      def matching_trains(trains)
+        trains.select { |id, _xml| id && @refs.include?(:train_numbers, id) }
+      end
+
+      def write_train_numbers(out, kept)
+        out << "<trainNumbers>"
+        kept.each { |pair| out << pair[1] }
+        out << "</trainNumbers>"
+      end
+
+      def read_train_number_children(reader)
+        depth = reader.depth
+        trains = []
+        return trains if reader.empty_element?
+
+        while reader.read
+          break if reader.node_type == TYPE_END && reader.depth == depth
+
+          if reader.node_type == TYPE_ELEMENT && reader.depth == depth + 1
+            if reader.local_name == "TrainNumber"
+              id = reader.attribute("id")
+              xml = reader.outer_xml
+              skip_element(reader)
+              trains << [id, xml] if id
+            else
+              skip_element(reader)
+            end
+          end
+        end
+
+        trains
+      end
+
+      def operator_only_keep?(op_ref)
+        return false if @line_filter
+
+        op_ref && @operator_ids.key?(op_ref)
+      end
+
+      def handle_text(reader, out)
+        value = reader.value.to_s
+        if pure_ws?(value)
+          buffer_ws(value)
+        else
+          flush_pending_ws(out)
+          out << value
+        end
+      end
+
+      def buffer_ws(value)
+        @pending_ws << value
+      end
+
+      def flush_pending_ws(out)
+        return if @pending_ws.empty?
+
+        out << @pending_ws
+        @pending_ws.clear
+      end
+
+      def clear_pending_ws
+        @pending_ws.clear
+      end
+
+      def pure_ws?(value)
+        value.match?(/\A\s*\z/)
+      end
+
+      def skip_to_depth(reader, depth)
+        while reader.read
+          break if reader.node_type == TYPE_END && reader.depth == depth
+        end
+      end
+
+      def append_leaf(reader, buf)
+        case reader.node_type
+        when TYPE_TEXT, TYPE_CDATA, TYPE_WS, TYPE_SWS
+          buf << reader.value.to_s
+        when TYPE_COMMENT
+          buf << "<!--#{reader.value}-->"
+        end
+      end
+
+      def write_start_to(reader, buf)
+        buf << "<#{reader.name}"
+        reader.attributes.each { |key, value| buf << " #{key}=\"#{esc_attr(value)}\"" }
+        buf << (reader.empty_element? ? "/>" : ">")
+      end
+
+      def skip_element(reader)
+        return if reader.empty_element?
+
+        depth = reader.depth
+        while reader.read
+          break if reader.node_type == TYPE_END && reader.depth == depth
+        end
+      end
+
+      def new_tempfile
+        file = Tempfile.new(["swiss-netex-tt", ".xml"])
+        file.binmode
+        file
+      end
+
+      def to_io(source)
+        return StringIO.new(source) if source.is_a?(String)
+        return source if source.respond_to?(:read)
+
+        raise Error, "Timetable expects a String or IO-like object"
+      end
+
+      def esc_attr(value)
+        value.to_s
+             .gsub("&", "&amp;")
+             .gsub('"', "&quot;")
+             .gsub("<", "&lt;")
+      end
+
+      def close_io(io)
+        return unless io
+
+        io.close!
+      rescue StandardError
+        io.close if io.respond_to?(:close) && !io.closed?
+      end
+    end
+  end
+end

data/lib/swiss_netex/filter/timetable_batch.rb

@@ -0,0 +1,330 @@
+# frozen_string_literal: true
+
+require "etc"
+require "json"
+
+module SwissNetex
+  class Filter
+    # Process TIMETABLE shards in short-lived child processes so libxml2/Nokogiri
+    # Reader heap is released at exit. Falls back to in-process when fork is
+    # unavailable (e.g. some Windows builds).
+    #
+    # Before forking, the parent byte-scans each shard for selected line /
+    # operator ids. Shards with no hits never parse as XML (big win on national
+    # drop-all shards up to ~600 MB uncompressed). Matching shards run in a
+    # process pool (SWISS_NETEX_TT_WORKERS, default: CPU count).
+    class TimetableBatch
+      # One shard per child keeps peak RSS ~size of the largest kept shard, not a
+      # multi-shard sum. Override with SWISS_NETEX_TT_BATCH.
+      DEFAULT_BATCH = 1
+      SCAN_CHUNK = 1 << 20 # 1 MiB
+      ShardResult = Data.define(:name, :basename, :journeys_kept)
+
+      def initialize(package_path:, tmp:, operator_ids:, line_ids:,
+                     line_filter: false, progress_io: nil, verbose: false)
+        @package_path = package_path.to_s
+        @tmp = tmp.to_s
+        @operator_ids = Array(operator_ids).map(&:to_s)
+        @line_ids = Array(line_ids).map(&:to_s)
+        @line_filter = line_filter
+        needles = (@line_ids + @operator_ids).reject(&:empty?).uniq
+        @max_needle = needles.map(&:bytesize).max || 0
+        @needle_re = needles.empty? ? nil : Regexp.union(needles)
+        @progress_io = progress_io
+        @verbose = verbose
+        @batch_size = resolve_batch_size
+        @workers = resolve_workers
+      end
+
+      def call(names, refs)
+        total = names.size
+        indexes = names.each_with_index.to_h { |name, i| [name, i + 1] }
+        candidates = select_candidates(names, indexes, total)
+
+        journey_count = 0
+        kept_by_index = {}
+
+        each_batch_result(candidates) do |results, delta|
+          refs.merge_export(delta)
+          results.each do |result|
+            index = indexes.fetch(result.name)
+            if result.journeys_kept.zero?
+              log "  [#{index}/#{total}] #{result.basename}: 0 journey(s)" if @verbose
+              next
+            end
+
+            journey_count += result.journeys_kept
+            kept_by_index[index] = result.basename
+            log "  [#{index}/#{total}] #{result.basename}: #{result.journeys_kept} journey(s)"
+          end
+        end
+
+        kept_names = kept_by_index.sort.map(&:last)
+        log "  journeys kept: #{journey_count} across #{kept_names.size} shard(s)"
+        [journey_count, kept_names]
+      end
+
+      private
+
+      def select_candidates(names, indexes, total)
+        return names if @needle_re.nil?
+
+        candidates = []
+        Package.open(@package_path) do |pkg|
+          names.each do |name|
+            pkg.open_member(name) do |io|
+              if may_match?(io)
+                candidates << name
+              elsif @verbose
+                index = indexes.fetch(name)
+                log "  [#{index}/#{total}] #{File.basename(name)}: skip (pre-scan)"
+              end
+            end
+          end
+        end
+
+        log "  pre-scan: #{candidates.size}/#{total} shard(s) may match " \
+            "(#{@workers} worker(s))"
+        candidates
+      end
+
+      def each_batch_result(names, &)
+        return if names.empty?
+
+        jobs = names.each_slice(@batch_size).to_a
+        if !fork_supported? || @workers <= 1 || jobs.size <= 1
+          jobs.each do |batch|
+            results, delta = run_batch(batch)
+            yield results, delta
+          end
+        else
+          run_pool(jobs, &)
+        end
+      end
+
+      def run_pool(jobs)
+        queue = jobs.dup
+        inflight = {}
+
+        while queue.any? || inflight.any?
+          while inflight.size < @workers && queue.any?
+            batch = queue.shift
+            spawn_job(batch, inflight)
+          end
+
+          pid, status = Process.wait2(-1)
+          meta = inflight.delete(pid)
+          unless meta
+            # Reap unexpected child rather than leave a zombie.
+            next
+          end
+
+          raw = meta[:reader].value
+          meta[:read_io].close
+          unless status.success?
+            reap_inflight(inflight)
+            raise Error, "TIMETABLE worker failed (exit #{status.exitstatus})"
+          end
+
+          results, delta = decode_payload(JSON.parse(raw))
+          yield results, delta
+        end
+      rescue JSON::ParserError => e
+        reap_inflight(inflight || {})
+        raise Error, "TIMETABLE worker returned invalid payload: #{e.message}"
+      rescue StandardError
+        reap_inflight(inflight || {})
+        raise
+      end
+
+      def spawn_job(batch, inflight)
+        read_io, write_io = IO.pipe
+        pid = fork do
+          read_io.close
+          payload = process_batch(batch, pre_scanned: true)
+          write_io.write(JSON.generate(encode_payload(payload)))
+          write_io.close
+          exit! 0
+        end
+        write_io.close
+        reader = Thread.new { read_io.read }
+        inflight[pid] = { read_io: read_io, reader: reader }
+      end
+
+      def reap_inflight(inflight)
+        inflight.each_key do |pid|
+          Process.kill("TERM", pid)
+        rescue Errno::ESRCH
+          # already gone
+        end
+        inflight.each do |pid, meta|
+          Process.wait2(pid)
+        rescue Errno::ECHILD
+          # already reaped
+        ensure
+          meta[:reader]&.kill
+          meta[:read_io]&.close unless meta[:read_io]&.closed?
+        end
+        inflight.clear
+      end
+
+      def run_batch(names)
+        return process_batch(names, pre_scanned: true) unless fork_supported?
+
+        read_io, write_io = IO.pipe
+        pid = fork do
+          read_io.close
+          payload = process_batch(names, pre_scanned: true)
+          write_io.write(JSON.generate(encode_payload(payload)))
+          write_io.close
+          exit! 0
+        end
+        write_io.close
+
+        raw = nil
+        status = nil
+        begin
+          raw = read_io.read
+        ensure
+          read_io.close
+          _pid, status = Process.wait2(pid)
+        end
+        raise Error, "TIMETABLE worker failed (exit #{status.exitstatus})" unless status.success?
+
+        decode_payload(JSON.parse(raw))
+      rescue JSON::ParserError => e
+        raise Error, "TIMETABLE worker returned invalid payload: #{e.message}"
+      end
+
+      def process_batch(names, pre_scanned: false)
+        local_refs = Refs.new
+        results = []
+
+        Package.open(@package_path) do |pkg|
+          names.each do |name|
+            results << process_shard(pkg, name, local_refs, pre_scanned: pre_scanned)
+          end
+        end
+
+        [results, local_refs.export]
+      end
+
+      def process_shard(pkg, name, local_refs, pre_scanned: false)
+        basename = File.basename(name)
+        journeys = 0
+
+        pkg.open_member(name) do |io|
+          unless pre_scanned || may_match?(io)
+            return ShardResult.new(name: name, basename: basename, journeys_kept: 0)
+          end
+
+          io.rewind if io.respond_to?(:rewind) && !pre_scanned
+          result = Timetable.new(
+            operator_ids: @operator_ids,
+            line_ids: @line_ids,
+            line_filter: @line_filter,
+            refs: local_refs
+          ).call(io)
+          begin
+            journeys = result.journeys_kept
+            copy_io_to_file(result.io, File.join(@tmp, basename)) if result.io
+          ensure
+            close_temp(result.io)
+          end
+        end
+
+        ShardResult.new(name: name, basename: basename, journeys_kept: journeys)
+      end
+
+      # Byte scan for selected ids. False positives only force a full Reader pass.
+      def may_match?(io)
+        return true unless @needle_re
+
+        tail = +""
+        while (chunk = io.read(SCAN_CHUNK))
+          haystack = tail.empty? ? chunk : (tail + chunk)
+          return true if haystack.match?(@needle_re)
+
+          tail = @max_needle.positive? ? haystack[-@max_needle, @max_needle].to_s : +""
+        end
+        false
+      end
+
+      def encode_payload(payload)
+        results, delta = payload
+        {
+          "results" => results.map do |r|
+            { "name" => r.name, "basename" => r.basename, "journeys_kept" => r.journeys_kept }
+          end,
+          "delta" => stringify_keys(delta)
+        }
+      end
+
+      def decode_payload(data)
+        results = data.fetch("results").map do |row|
+          ShardResult.new(
+            name: row.fetch("name"),
+            basename: row.fetch("basename"),
+            journeys_kept: row.fetch("journeys_kept")
+          )
+        end
+        delta = symbolize_keys(data.fetch("delta"))
+        [results, delta]
+      end
+
+      def stringify_keys(hash)
+        hash.transform_keys(&:to_s).transform_values { |ids| Array(ids).map(&:to_s) }
+      end
+
+      def symbolize_keys(hash)
+        hash.transform_keys(&:to_sym)
+      end
+
+      def fork_supported?
+        return false if ENV["SWISS_NETEX_NO_FORK"]
+
+        Process.respond_to?(:fork)
+      end
+
+      def resolve_batch_size
+        raw = ENV["SWISS_NETEX_TT_BATCH"].to_s.strip
+        return DEFAULT_BATCH if raw.empty?
+
+        size = Integer(raw, 10)
+        size.positive? ? size : DEFAULT_BATCH
+      rescue ArgumentError
+        DEFAULT_BATCH
+      end
+
+      def resolve_workers
+        raw = ENV["SWISS_NETEX_TT_WORKERS"].to_s.strip
+        if raw.empty?
+          n = Etc.nprocessors
+          return n.positive? ? n : 1
+        end
+
+        size = Integer(raw, 10)
+        size.positive? ? size : 1
+      rescue ArgumentError
+        1
+      end
+
+      def copy_io_to_file(io, destination)
+        io.rewind if io.respond_to?(:rewind)
+        File.open(destination, "wb") { |out| IO.copy_stream(io, out) }
+      end
+
+      def close_temp(io)
+        return unless io
+
+        io.close!
+      rescue StandardError
+        io.close if io.respond_to?(:close) && !io.closed?
+      end
+
+      def log(message)
+        @progress_io&.puts(message)
+      end
+    end
+  end
+end