6.4 增量式解析大型XML文件

問題

你想使用盡可能少的內(nèi)存從一個超大的XML文檔中提取數(shù)據(jù)。

解決方案

任何時候只要你遇到增量式的數(shù)據(jù)處理時，第一時間就應該想到迭代器和生成器。下面是一個很簡單的函數(shù)，只使用很少的內(nèi)存就能增量式的處理一個大型XML文件：

from xml.etree.ElementTree import iterparse

def parse_and_remove(filename, path):
    path_parts = path.split('/')
    doc = iterparse(filename, ('start', 'end'))
    # Skip the root element
    next(doc)

    tag_stack = []
    elem_stack = []
    for event, elem in doc:
        if event == 'start':
            tag_stack.append(elem.tag)
            elem_stack.append(elem)
        elif event == 'end':
            if tag_stack == path_parts:
                yield elem
                elem_stack[-2].remove(elem)
            try:
                tag_stack.pop()
                elem_stack.pop()
            except IndexError:
                pass

為了測試這個函數(shù)，你需要先有一個大型的XML文件。通常你可以在政府網(wǎng)站或公共數(shù)據(jù)網(wǎng)站上找到這樣的文件。例如，你可以下載XML格式的芝加哥城市道路坑洼數(shù)據(jù)庫。在寫這本書的時候，下載文件已經(jīng)包含超過100,000行數(shù)據(jù)，編碼格式類似于下面這樣：

假設你想寫一個腳本來按照坑洼報告數(shù)量排列郵編號碼。你可以像這樣做：

from xml.etree.ElementTree import parse
from collections import Counter

potholes_by_zip = Counter()

doc = parse('potholes.xml')
for pothole in doc.iterfind('row/row'):
    potholes_by_zip[pothole.findtext('zip')] += 1
for zipcode, num in potholes_by_zip.most_common():
    print(zipcode, num)

這個腳本唯一的問題是它會先將整個XML文件加載到內(nèi)存中然后解析。在我的機器上，為了運行這個程序需要用到450MB左右的內(nèi)存空間。如果使用如下代碼，程序只需要修改一點點：

from collections import Counter

potholes_by_zip = Counter()

data = parse_and_remove('potholes.xml', 'row/row')
for pothole in data:
    potholes_by_zip[pothole.findtext('zip')] += 1
for zipcode, num in potholes_by_zip.most_common():
    print(zipcode, num)

結(jié)果是：這個版本的代碼運行時只需要7MB的內(nèi)存–大大節(jié)約了內(nèi)存資源。

討論

這一節(jié)的技術(shù)會依賴 ElementTree 模塊中的兩個核心功能。第一，iterparse() 方法允許對XML文檔進行增量操作。使用時，你需要提供文件名和一個包含下面一種或多種類型的事件列表：start , end, start-ns 和 end-ns 。由 iterparse() 創(chuàng)建的迭代器會產(chǎn)生形如 (event, elem) 的元組，其中 event 是上述事件列表中的某一個，而 elem 是相應的XML元素。例如：

>>> data = iterparse('potholes.xml',('start','end'))
>>> next(data)
('start', <Element 'response' at 0x100771d60>)
>>> next(data)
('start', <Element 'row' at 0x100771e68>)
>>> next(data)
('start', <Element 'row' at 0x100771fc8>)
>>> next(data)
('start', <Element 'creation_date' at 0x100771f18>)
>>> next(data)
('end', <Element 'creation_date' at 0x100771f18>)
>>> next(data)
('start', <Element 'status' at 0x1006a7f18>)
>>> next(data)
('end', <Element 'status' at 0x1006a7f18>)
>>>

start 事件在某個元素第一次被創(chuàng)建并且還沒有被插入其他數(shù)據(jù)(如子元素)時被創(chuàng)建。而 end 事件在某個元素已經(jīng)完成時被創(chuàng)建。盡管沒有在例子中演示，start-ns 和 end-ns 事件被用來處理XML文檔命名空間的聲明。

這本節(jié)例子中，start 和 end 事件被用來管理元素和標簽棧。棧代表了文檔被解析時的層次結(jié)構(gòu)，還被用來判斷某個元素是否匹配傳給函數(shù) parse_and_remove() 的路徑。如果匹配，就利用 yield 語句向調(diào)用者返回這個元素。

在 yield 之后的下面這個語句才是使得程序占用極少內(nèi)存的ElementTree的核心特性：

elem_stack[-2].remove(elem)

這個語句使得之前由 yield 產(chǎn)生的元素從它的父節(jié)點中刪除掉。假設已經(jīng)沒有其它的地方引用這個元素了，那么這個元素就被銷毀并回收內(nèi)存。

對節(jié)點的迭代式解析和刪除的最終效果就是一個在文檔上高效的增量式清掃過程。文檔樹結(jié)構(gòu)從始自終沒被完整的創(chuàng)建過。盡管如此，還是能通過上述簡單的方式來處理這個XML數(shù)據(jù)。

這種方案的主要缺陷就是它的運行性能了。我自己測試的結(jié)果是，讀取整個文檔到內(nèi)存中的版本的運行速度差不多是增量式處理版本的兩倍快。但是它卻使用了超過后者60倍的內(nèi)存。因此，如果你更關(guān)心內(nèi)存使用量的話，那么增量式的版本完勝。