#!/bin/env -S poetry run pythonfromconcurrent.futuresimportThreadPoolExecutor, as_completedimportpyarrowaspaimportpyarrow.datasetaspdfromadlfsimportAzureBlobFileSystemfromazure.identityimportDefaultAzureCredentialfromazure.storage.blobimportBlobServiceClientfrompyarrow.datasetimportHivePartitioningfrompyiceberg.catalogimportCatalogfrompyiceberg.catalog.sqlimportSqlCatalogfrompyiceberg.io.pyarrowimportpyarrow_to_schemafrompyiceberg.partitioningimportPartitionField, PartitionSpecfrompyiceberg.table.name_mappingimportMappedField, NameMappingfrompyiceberg.transformsimportIdentityTransformimportsettingsclassAzureStorage:
def__init__(self):
credential=DefaultAzureCredential()
blob_service_client=BlobServiceClient(
settings.AZURE_BLOB_URL, credential
)
self.container_client=blob_service_client.get_container_client(
settings.AZURE_BLOB_CONTAINER
)
# The AzureBlobFileSystem doesn't cleanly shutdown and currently# always raises an expection at the end of this program. See:# https://github.com/fsspec/adlfs/issues/431self.abfs=AzureBlobFileSystem(
account_name=settings.AZURE_BLOB_ACCOUNT_NAME,
credential=credential,
)
deflist_tables(self):
returnself.container_client.walk_blobs(
settings.AZURE_LIVE_PATH, delimiter="/"
)
defload_dataset(self, table_name) ->pd.Dataset:
name="/".join((settings.AZURE_LIVE_PATH.rstrip("/"), table_name))
dataset=pd.dataset(
"/".join([settings.AZURE_LIVE_CONTAINER, name]),
format="parquet",
filesystem=self.abfs,
partitioning=HivePartitioning(
pa.schema(
[
("dataset", pa.string()),
("flavor", pa.string()),
]
)
),
)
returndatasetdefcreate_iceberg_catalog():
catalog=SqlCatalog(
"default",
**{
"uri": settings.ICEBERG_DATABASE_URI,
"warehouse": settings.ICEBERG_WAREHOUSE,
},
)
returncatalogdefdownload_table(catalog: Catalog, table_name: str, ds: pd.Dataset):
name_mapping=NameMapping(
root=[
MappedField(field_id=field_id, names=[field.name])
forfield_id, fieldinenumerate(ds.schema, 1)
]
)
schema=pyarrow_to_schema(ds.schema, name_mapping=name_mapping)
assertisinstance(ds.partitioning, HivePartitioning), ds.partitioningpartitioning_spec=PartitionSpec(
*(
PartitionField(
source_id=name_mapping.find(field.name).field_id,
field_id=-1,
transform=IdentityTransform(),
name=field.name,
)
forfieldinds.partitioning.schema
)
)
table=catalog.create_table(
f"{settings.ICEBERG_NAMESPACE}.{table_name}",
schema=schema,
partition_spec=partitioning_spec,
)
fragments=list(ds.get_fragments())
withThreadPoolExecutor(8) asexecutor:
futures= [
executor.submit(
download_fragment,
table.identifier,
fragment,
)
forfragmentinfragments
]
forfutureinas_completed(futures):
try:
future.result()
exceptExceptionase:
executor.shutdown(wait=False, cancel_futures=True)
raiseefromNonedefdownload_fragment(
table_identifier: str,
fragment,
):
catalog=create_iceberg_catalog()
partition_keys=pd.get_partition_keys(fragment.partition_expression)
fragment_table=fragment.to_table()
fork, vinpartition_keys.items():
fragment_table=fragment_table.append_column(
pa.field(k, pa.string(), nullable=False),
pa.repeat(pa.scalar(v), fragment_table.num_rows),
)
table=catalog.load_table(table_identifier)
table.append(fragment_table)
defimport_data(storage: AzureStorage, catalog, table_name):
dataset=storage.load_dataset(table_name)
download_table(catalog, table_name, dataset)
defmain():
catalog=create_iceberg_catalog()
catalog.create_namespace_if_not_exists(settings.ICEBERG_NAMESPACE)
storage=AzureStorage()
fortable_nameinstorage.list_tables():
import_data(storage, catalog, table_name)
if__name__=="__main__":
main()
Apache Iceberg version
0.7.1
Please describe the bug 🐞
Summary
I'm currently trying to migrate a couple of dataframes with a custom hive-like storage scheme to Iceberg. After a lot of fiddling I managed to load the dataframes from an Azure storage, create the table in the Iceberg catalog (currently using sqlite + local fs) and append fragments from the Parquet dataset. As soon as adding a thread pool I always run into concurrency issues.
Errors
I get either of the following two error messages:
or
Sources
I use
Dataset.get_fragmentsand insert the data into an iceberg table with identical partitioning.I can work around this error by using a GIL (global iceberg lock, pun intended.) which is just a
threading.Lock()that ensures everyload_table()+table.appendhappens atomically. But that kills almost all performance gains there could be made. Also I plan on using this in some Celery runners . So using athreading.Lock()is no option in the future anyways.azure_import.py
pyproject.toml