(
table: Table,
data_source: DataSource,
repo_path: str,
mode: str = "append",
allow_overwrite: bool = False,
)
| 58 | |
| 59 | |
| 60 | def _write_data_source( |
| 61 | table: Table, |
| 62 | data_source: DataSource, |
| 63 | repo_path: str, |
| 64 | mode: str = "append", |
| 65 | allow_overwrite: bool = False, |
| 66 | ): |
| 67 | assert isinstance(data_source, FileSource) |
| 68 | |
| 69 | file_options = data_source.file_options |
| 70 | |
| 71 | absolute_path = FileSource.get_uri_for_file_path( |
| 72 | repo_path=repo_path, uri=file_options.uri |
| 73 | ) |
| 74 | |
| 75 | if ( |
| 76 | mode == "overwrite" |
| 77 | and not allow_overwrite |
| 78 | and os.path.exists(str(absolute_path)) |
| 79 | ): |
| 80 | raise SavedDatasetLocationAlreadyExists(location=file_options.uri) |
| 81 | |
| 82 | if data_source.file_format is None or isinstance( |
| 83 | data_source.file_format, ParquetFormat |
| 84 | ): |
| 85 | if mode == "overwrite": |
| 86 | table = table.to_pyarrow() |
| 87 | |
| 88 | filesystem, path = FileSource.create_filesystem_and_path( |
| 89 | str(absolute_path), |
| 90 | file_options.s3_endpoint_override, |
| 91 | ) |
| 92 | |
| 93 | if path.endswith(".parquet"): |
| 94 | pyarrow.parquet.write_table(table, where=path, filesystem=filesystem) |
| 95 | else: |
| 96 | # otherwise assume destination is directory |
| 97 | pyarrow.parquet.write_to_dataset( |
| 98 | table, root_path=path, filesystem=filesystem |
| 99 | ) |
| 100 | elif mode == "append": |
| 101 | table = table.to_pyarrow() |
| 102 | prev_table = ibis.read_parquet(file_options.uri).to_pyarrow() |
| 103 | if table.schema != prev_table.schema: |
| 104 | table = table.cast(prev_table.schema) |
| 105 | new_table = pyarrow.concat_tables([table, prev_table]) |
| 106 | ibis.memtable(new_table).to_parquet(file_options.uri) |
| 107 | elif isinstance(data_source.file_format, DeltaFormat): |
| 108 | storage_options = { |
| 109 | "AWS_ENDPOINT_URL": str(data_source.s3_endpoint_override), |
| 110 | } |
| 111 | |
| 112 | if mode == "append": |
| 113 | from deltalake import DeltaTable |
| 114 | |
| 115 | prev_schema = ( |
| 116 | DeltaTable(file_options.uri, storage_options=storage_options) |
| 117 | .schema() |
nothing calls this directly
no test coverage detected