Zenn
🐻‍❄️

polarsでCSVファイルを読み書きする

2023/10/28に公開
Polars
tech

CSVファイルの読み書き

  • df = pl.read_csv(filename): 読み込み
  • df.write_csv(filename): 書き出し

遅延データフレームへの読み書き

  • lf = pl.scan_csv(filename): 遅延読み込み
  • lf.sink_csv(filename): 遅延データフレームを書き出し

その他のファイルの読み書き

CSVファイルと同様に、excelスプレッドシート/JSONファイル/データベースなど様々なファイルを読み書きできる。

  • pl.read_excel(filename), df.write_excel(filename)
  • pl.read_json(filename), df.write_json(filename)
  • pl.read_ndjson(filename), df.write_ndjson(filename): rootが複数あるJSONファイルを読み書きできて便利
  • pl.read_database(filename), df.write_database(filename)
JSONファイルの形式

read_jsonで読み込む(書き出す)ことは以下のファイルで可能。

sample.json
{
    "columns":[
        {
            "name":"name",
            "datatype":"Utf8",
            "bit_settings":"",
            "values":[
                "natori","kenmochi","iincho","lize","shiina","chaika","sigureui","udukohh","ririmu","hoshikawa"
                ]
        },
        {
            "name":"column1",
            "datatype":"Utf8",
            "bit_settings":"",
            "values":["38"," ","35","58","45","47","38","36","56","34"]
        },
        {
            "name":"column2",
            "datatype":"Int64",
            "bit_settings":"",
            "values":[0,1,1,0,0,1,1,0,0,1]},
        {
            "name":"column3",
            "datatype":"Int64",
            "bit_settings":"",
            "values":[
                6147,3192,6463,1733,796,1319,1312,3400,1594,8722
            ]
        },
        {
            "name":"column4",
            "datatype":"Utf8",
            "bit_settings":"",
            "values":[
                " ","5969","3773","4413","7541","1449","6031","1363","3632","5961"
            ]
        },
        {
            "name":"column5",
            "datatype":"Int64",
            "bit_settings":"",
            "values":[
                2471,5878,832,2252,3098,4234,9885,1427,9449,8581
            ]
        }
    ]
}

shape: (10, 6)
┌───────────┬─────────┬─────────┬─────────┬─────────┬─────────┐
│ name      ┆ column1 ┆ column2 ┆ column3 ┆ column4 ┆ column5 │
│ ---       ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
│ str       ┆ str     ┆ i64     ┆ i64     ┆ str     ┆ i64     │
╞═══════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
│ natori    ┆ 38      ┆ 0       ┆ 6147    ┆         ┆ 2471    │
│ kenmochi  ┆         ┆ 1       ┆ 3192    ┆ 5969    ┆ 5878    │
│ iincho    ┆ 35      ┆ 1       ┆ 6463    ┆ 3773    ┆ 832     │
│ lize      ┆ 58      ┆ 0       ┆ 1733    ┆ 4413    ┆ 2252    │
│ …         ┆ …       ┆ …       ┆ …       ┆ …       ┆ …       │
│ sigureui  ┆ 38      ┆ 1       ┆ 1312    ┆ 6031    ┆ 9885    │
│ udukohh   ┆ 36      ┆ 0       ┆ 3400    ┆ 1363    ┆ 1427    │
│ ririmu    ┆ 56      ┆ 0       ┆ 1594    ┆ 3632    ┆ 9449    │
│ hoshikawa ┆ 34      ┆ 1       ┆ 8722    ┆ 5961    ┆ 8581    │
└───────────┴─────────┴─────────┴─────────┴─────────┴─────────┘

https://pola-rs.github.io/polars/py-polars/html/reference/io.html

https://pola-rs.github.io/polars/user-guide/io/csv/

CSVの読み込み

sample.csv
name,column1,column2,column3,column4,column5
natori,38,0,6147,1164,2471
kenmochi,46,1,3192,5969,5878
iincho,35,1,6463,3773,832
lize,58,0,1733,4413,2252
shiina,45,0,796,7541,3098
chaika,47,1,1319,1449,4234
sigureui,38,1,1312,6031,9885
udukohh,36,0,3400,1363,1427
ririmu,56,0,1594,3632,9449
hoshikawa,34,1,8722,5961,8581

これをデータフレームに読み込むには次のようにする。

df = pl.read_csv("sample.csv")
print(df)
# shape: (10, 6)
# ┌───────────┬─────────┬─────────┬─────────┬─────────┬─────────┐
# │ name      ┆ column1 ┆ column2 ┆ column3 ┆ column4 ┆ column5 │
# │ ---       ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
# │ str       ┆ i64     ┆ i64     ┆ i64     ┆ i64     ┆ i64     │
# ╞═══════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
# │ natori    ┆ 38      ┆ 0       ┆ 6147    ┆ 1164    ┆ 2471    │
# │ kenmochi  ┆ 46      ┆ 1       ┆ 3192    ┆ 5969    ┆ 5878    │
# │ iincho    ┆ 35      ┆ 1       ┆ 6463    ┆ 3773    ┆ 832     │
# │ lize      ┆ 58      ┆ 0       ┆ 1733    ┆ 4413    ┆ 2252    │
# │ …         ┆ …       ┆ …       ┆ …       ┆ …       ┆ …       │
# │ hoshikawa ┆ 34      ┆ 1       ┆ 8722    ┆ 5961    ┆ 8581    │
# └───────────┴─────────┴─────────┴─────────┴─────────┴─────────┘

print(df.columns)
# ['name', 'column1', 'column2', 'column3', 'column4', 'column5']

read_csvのoption引数で、主に使いそうなのは以下。

  • columns: Sequence[int] | Sequence[str] | None = None

  • new_columns: Sequence[str] | None = None

  • dtypes: Mapping[str, PolarsDataType] | Sequence[PolarsDataType] | None = None

  • has_header: bool = True ←CSVに列名が存在するかどうか

  • encoding: str = "utf8"

  • separator: str = ","

  • quote_char: str | None = '"'

  • eol_char: str = "\n"

  • null_values: str | Sequence[str] | dict[str str] | None = None

  • n_rows: int | None = None

  • skip_rows: int = 0

特定の列のみを読み込み

columnsに読み込む列名のリストを与えると、その列だけが読み込まれる。

df = pl.read_csv("sample.csv", columns=['name', 'column1', 'column2'])
print(df)
# shape: (10, 3)
# ┌───────────┬─────────┬─────────┐
# │ name      ┆ column1 ┆ column2 │
# │ ---       ┆ ---     ┆ ---     │
# │ str       ┆ i64     ┆ i64     │
# ╞═══════════╪═════════╪═════════╡
# │ natori    ┆ 38      ┆ 0       │
# │ kenmochi  ┆ 46      ┆ 1       │
# │ iincho    ┆ 35      ┆ 1       │
# │ lize      ┆ 58      ┆ 0       │
# │ …         ┆ …       ┆ …       │
# │ hoshikawa ┆ 34      ┆ 1       │
# └───────────┴─────────┴─────────┘

列名を貼り替えてread

new_columnsに新しくつける列名のリストを与える。
このとき古い列名は捨てられるので注意。

df = pl.read_csv("sample.csv", new_columns=['n', 'c1', 'c2', 'c3', 'c4', 'c5'])
print(df)
# shape: (10, 6)
# ┌───────────┬─────┬─────┬──────┬──────┬──────┐
# │ n         ┆ c1  ┆ c2  ┆ c3   ┆ c4   ┆ c5   │
# │ ---       ┆ --- ┆ --- ┆ ---  ┆ ---  ┆ ---  │
# │ str       ┆ i64 ┆ i64 ┆ i64  ┆ i64  ┆ i64  │
# ╞═══════════╪═════╪═════╪══════╪══════╪══════╡
# │ natori    ┆ 38  ┆ 0   ┆ 6147 ┆ 1164 ┆ 2471 │
# │ kenmochi  ┆ 46  ┆ 1   ┆ 3192 ┆ 5969 ┆ 5878 │
# │ iincho    ┆ 35  ┆ 1   ┆ 6463 ┆ 3773 ┆ 832  │
# │ lize      ┆ 58  ┆ 0   ┆ 1733 ┆ 4413 ┆ 2252 │
# │ …         ┆ …   ┆ …   ┆ …    ┆ …    ┆ …    
# │ hoshikawa ┆ 34  ┆ 1   ┆ 8722 ┆ 5961 ┆ 8581 │
# └───────────┴─────┴─────┴──────┴──────┴──────┘

データ型を指定

dtypesへ列名とその型を示すdictを与える。
型はPythonネイティブの型ではなくpl.DataTypeであることに注意。

https://pola-rs.github.io/polars/py-polars/html/reference/datatypes.html

f32 = pl.Float32
df = pl.read_csv(
    "sample.csv", 
    dtypes={"name": str, "column1": f32, "column2": f32, "column3": f32, "column4": f32, "column5": f32}
    )
print(df)
# shape: (10, 6)
# ┌───────────┬─────────┬─────────┬─────────┬─────────┬─────────┐
# │ name      ┆ column1 ┆ column2 ┆ column3 ┆ column4 ┆ column5 │
# │ ---       ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
# │ str       ┆ f32     ┆ f32     ┆ f32     ┆ f32     ┆ f32     │
# ╞═══════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
# │ natori    ┆ 38.0    ┆ 0.0     ┆ 6147.0  ┆ 1164.0  ┆ 2471.0  │
# │ kenmochi  ┆ 46.0    ┆ 1.0     ┆ 3192.0  ┆ 5969.0  ┆ 5878.0  │
# │ iincho    ┆ 35.0    ┆ 1.0     ┆ 6463.0  ┆ 3773.0  ┆ 832.0   │
# │ lize      ┆ 58.0    ┆ 0.0     ┆ 1733.0  ┆ 4413.0  ┆ 2252.0  │
# │ …         ┆ …       ┆ …       ┆ …       ┆ …       ┆ …       │
# │ hoshikawa ┆ 34.0    ┆ 1.0     ┆ 8722.0  ┆ 5961.0  ┆ 8581.0  │
# └───────────┴─────────┴─────────┴─────────┴─────────┴─────────┘

指定の行数分読み込む

n_rowsに与えた値だけ先頭から読み取る。

df = pl.read_csv("sample.csv", n_rows=4)
print(df)
# shape: (4, 6)
# ┌──────────┬─────────┬─────────┬─────────┬─────────┬─────────┐
# │ name     ┆ column1 ┆ column2 ┆ column3 ┆ column4 ┆ column5 │
# │ ---      ┆ ---     ┆ ---     ┆ ---     ┆ ---     ┆ ---     │
# │ str      ┆ i64     ┆ i64     ┆ i64     ┆ i64     ┆ i64     │
# ╞══════════╪═════════╪═════════╪═════════╪═════════╪═════════╡
# │ natori   ┆ 38      ┆ 0       ┆ 6147    ┆ 1164    ┆ 2471    │
# │ kenmochi ┆ 46      ┆ 1       ┆ 3192    ┆ 5969    ┆ 5878    │
# │ iincho   ┆ 35      ┆ 1       ┆ 6463    ┆ 3773    ┆ 832     │
# │ lize     ┆ 58      ┆ 0       ┆ 1733    ┆ 4413    ┆ 2252    │
# └──────────┴─────────┴─────────┴─────────┴─────────┴─────────┘

CSVファイルの先頭をスキップ

skip_rowsで先頭を飛ばせる。ただし列名に注意。

df = pl.read_csv("sample.csv", skip_rows=5)
print(df)
# shape: (5, 6)
# ┌───────────┬─────┬─────┬──────┬──────┬──────┐
# │ shiina    ┆ 45  ┆ 0   ┆ 796  ┆ 7541 ┆ 3098 │
# │ ---       ┆ --- ┆ --- ┆ ---  ┆ ---  ┆ ---  │
# │ str       ┆ i64 ┆ i64 ┆ i64  ┆ i64  ┆ i64  │
# ╞═══════════╪═════╪═════╪══════╪══════╪══════╡
# │ chaika    ┆ 47  ┆ 1   ┆ 1319 ┆ 1449 ┆ 4234 │
# │ sigureui  ┆ 38  ┆ 1   ┆ 1312 ┆ 6031 ┆ 9885 │
# │ udukohh   ┆ 36  ┆ 0   ┆ 3400 ┆ 1363 ┆ 1427 │
# │ ririmu    ┆ 56  ┆ 0   ┆ 1594 ┆ 3632 ┆ 9449 │
# │ hoshikawa ┆ 34  ┆ 1   ┆ 8722 ┆ 5961 ┆ 8581 │
# └───────────┴─────┴─────┴──────┴──────┴──────┘

列名無しのCSV

has_headerをFalseにするとCSVの先頭から値として読み込む。
列名はデフォルトでcolumn_nになるが、new_columnsを与えることで命名して読み込める。

df = pl.read_csv("sample.csv", has_header=False)
# shape: (10, 6)
# ┌───────────┬──────────┬──────────┬──────────┬──────────┬──────────┐
# │ column_1  ┆ column_2 ┆ column_3 ┆ column_4 ┆ column_5 ┆ column_6 │
# │ ---       ┆ ---      ┆ ---      ┆ ---      ┆ ---      ┆ ---      │
# │ str       ┆ str      ┆ i64      ┆ i64      ┆ str      ┆ i64      │
# ╞═══════════╪══════════╪══════════╪══════════╪══════════╪══════════╡
# │ natori    ┆ 38       ┆ 0        ┆ 6147     ┆          ┆ 2471     │
# │ kenmochi  ┆          ┆ 1        ┆ 3192     ┆ 5969     ┆ 5878     │
# │ iincho    ┆ 35       ┆ 1        ┆ 6463     ┆ 3773     ┆ 832      │
# │ lize      ┆ 58       ┆ 0        ┆ 1733     ┆ 4413     ┆ 2252     │
# │ …         ┆ …        ┆ …        ┆ …        ┆ …        ┆ …        │
# │ hoshikawa ┆ 34       ┆ 1        ┆ 8722     ┆ 5961     ┆ 8581     │
# └───────────┴──────────┴──────────┴──────────┴──────────┴──────────┘

CSVの遅延読み込み

scan_csvを使うと遅延評価モードのデータフレーム(LazyFrame)を作れる。

df = pl.scan_csv("sample.csv")
print(df)

# naive plan: (run LazyFrame.explain(optimized=True) to see the optimized plan)
#
#   Csv SCAN sample.csv
#   PROJECT */6 COLUMNS

Discussion