Datasets documentation
PyArrow와 함께 사용
PyArrow와 함께 사용
이 문서는 datasets를 PyArrow와 함께 사용하는 방법에 대한 간단한 소개입니다. 특히 Arrow compute 함수를 사용하여 데이터를 처리하는 방법과, 데이터를 PyArrow로 변환하거나 PyArrow에서 가져오는 방법에 중점을 둡니다.
datasets는 내부적으로 PyArrow를 사용하므로 제로-복사(zero-copy) 연산을 빠르게 수행할 수 있어 특히 유용합니다.
데이터셋 형식
기본적으로 datasets는 정수, 실수, 문자열, 리스트 등 일반적인 Python 객체를 반환합니다.
대신 PyArrow 테이블 또는 배열을 얻으려면 Dataset.with_format()을 사용하여 데이터셋의 형식을 pyarrow로 설정할 수 있습니다:
>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("arrow")
>>> ds[0] # pa.Table
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a"]]
col_1: [[0]]
>>> ds[:2] # pa.Table
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a","b"]]
col_1: [[0,0]]
>>> ds["data"] # pa.array
<pyarrow.lib.ChunkedArray object at 0x1394312a0>
[
[
"a",
"b",
"c",
"d"
]
]IterableDataset 객체, 예를 들어 load_dataset(..., streaming=True)를 사용하여 가져온 객체에도 적용됩니다:
>>> ds = ds.with_format("arrow")
>>> for table in ds.iter(batch_size=2):
... print(table)
... break
pyarrow.Table
col_0: string
col_1: double
----
col_0: [["a","b"]]
col_1: [[0,0]]데이터 처리
PyArrow 함수는 일반적으로 직접 작성한 Python 함수보다 빠르므로 데이터 처리를 최적화하는 좋은 선택지입니다. Arrow compute 함수를 사용하여 Dataset.map() 또는 Dataset.filter()에서 데이터셋을 처리할 수 있습니다:
>>> import pyarrow.compute as pc
>>> from datasets import Dataset
>>> data = {"col_0": ["a", "b", "c", "d"], "col_1": [0., 0., 1., 1.]}
>>> ds = Dataset.from_dict(data)
>>> ds = ds.with_format("arrow")
>>> ds = ds.map(lambda t: t.append_column("col_2", pc.add(t["col_1"], 1)), batched=True)
>>> ds[:2]
pyarrow.Table
col_0: string
col_1: double
col_2: double
----
col_0: [["a","b"]]
col_1: [[0,0]]
col_2: [[1,1]]
>>> ds = ds.filter(lambda t: pc.equal(t["col_0"], "b"), batched=True)
>>> ds[0]
pyarrow.Table
col_0: string
col_1: double
col_2: double
----
col_0: [["b"]]
col_1: [[0]]
col_2: [[1]]PyArrow에서 데이터를 행별로 처리하는 것보다 배치 단위로 처리하는 것이 더 빠르기 때문에 batched=True를 사용합니다. 또한 map()에서 batch_size=를 사용하여 각 table의 크기를 설정할 수도 있습니다.
이것은 IterableDataset.map()과 IterableDataset.filter()에도 적용됩니다.
PyArrow에서 가져오기 또는 내보내기
Dataset은 PyArrow Table의 래퍼입니다. Table에서 Dataset을 직접 생성할 수 있습니다:
ds = Dataset(table)
Dataset.data를 사용하면 데이터셋의 PyArrow Table에 접근할 수 있으며, Arrow 데이터의 출처와 적용된 연산에 따라 MemoryMappedTable 또는 InMemoryTable 또는 ConcatenationTable을 반환합니다.
이러한 객체는 Dataset.data.table에서 접근할 수 있는 기본 PyArrow 테이블을 래핑합니다. 이 테이블에는 데이터셋의 모든 데이터가 포함되어 있지만, Dataset._indices에 데이터셋 행 인덱스를 PyArrow Table 행 인덱스로 매핑하는 인덱스 매핑이 있을 수도 있습니다. 이는 Dataset.shuffle()로 데이터셋이 셔플되었거나, 행의 일부만 사용된 경우(예: Dataset.select() 이후)에 발생할 수 있습니다.
일반적인 경우, table = ds.with_format("arrow")[:]를 사용하여 데이터셋을 PyArrow Table로 내보낼 수 있습니다.