read_sql 成塊的 polars

Question

我正在嘗試讀取一個帶有極坐標的大型數據庫表。 不幸的是，數據太大無法放入 memory，下面的代碼最終失敗了。

polars 中有沒有一種方法可以定義塊大小，並將這些塊寫入鑲木地板，或者使用惰性 dataframe 接口來保持 memory 占用空間低？

import polars as pl
df = pl.read_sql("SELECT * from TABLENAME", connection_string) 
df.write_parquet("output.parquet")

Answer 1

是和不是。

沒有預定義的方法來執行此操作，但您當然可以自己執行。 你會做這樣的事情：

rows_at_a_time=1000
curindx=0
while True:
    df = pl.read_sql(f"SELECT * from TABLENAME limit {curindx},{rows_at_a_time}", connection_string) 
    if df.shape[0]==0:
        break
    df.write_parquet(f"output{curindx}.parquet")
    curindx+=rows_at_a_time
ldf=pl.concat([pl.scan_parquet(x) for x in os.listdir(".") if "output" in x and "parquet" in x])

假設您使用的是 mysql 或具有相同語法的數據庫（這不是微不足道的假設），這會從該答案中借用限制語法。 如果不使用 mysql，您可能需要執行類似的操作。

否則，您只需分塊讀取表，將每個塊保存到本地文件中。 當您從查詢中返回的塊有 0 行時，它會停止循環並將所有文件加載到惰性 df。

您幾乎可以肯定（並且應該）將rows_at_a_time增加到大於 1000 的值，但這取決於您的數據和計算機 memory。