| 1227 | |
| 1228 | |
| 1229 | def _add_any_metadata(table, pandas_metadata): |
| 1230 | modified_columns = {} |
| 1231 | modified_fields = {} |
| 1232 | |
| 1233 | schema = table.schema |
| 1234 | |
| 1235 | index_columns = pandas_metadata['index_columns'] |
| 1236 | # only take index columns into account if they are an actual table column |
| 1237 | index_columns = [idx_col for idx_col in index_columns |
| 1238 | if isinstance(idx_col, str)] |
| 1239 | n_index_levels = len(index_columns) |
| 1240 | n_columns = len(pandas_metadata['columns']) - n_index_levels |
| 1241 | |
| 1242 | # Add time zones |
| 1243 | for i, col_meta in enumerate(pandas_metadata['columns']): |
| 1244 | |
| 1245 | raw_name = col_meta.get('field_name') |
| 1246 | if not raw_name: |
| 1247 | # deal with metadata written with arrow < 0.8 or fastparquet |
| 1248 | raw_name = col_meta['name'] |
| 1249 | if i >= n_columns: |
| 1250 | # index columns |
| 1251 | raw_name = index_columns[i - n_columns] |
| 1252 | if raw_name is None: |
| 1253 | raw_name = 'None' |
| 1254 | |
| 1255 | idx = schema.get_field_index(raw_name) |
| 1256 | if idx != -1: |
| 1257 | if col_meta['pandas_type'] == 'datetimetz': |
| 1258 | col = table[idx] |
| 1259 | if not isinstance(col.type, pa.lib.TimestampType): |
| 1260 | continue |
| 1261 | metadata = col_meta['metadata'] |
| 1262 | if not metadata: |
| 1263 | continue |
| 1264 | metadata_tz = metadata.get('timezone') |
| 1265 | if metadata_tz and metadata_tz != col.type.tz: |
| 1266 | converted = col.to_pandas() |
| 1267 | tz_aware_type = pa.timestamp('ns', tz=metadata_tz) |
| 1268 | with_metadata = pa.Array.from_pandas(converted, |
| 1269 | type=tz_aware_type) |
| 1270 | |
| 1271 | modified_fields[idx] = pa.field(schema[idx].name, |
| 1272 | tz_aware_type) |
| 1273 | modified_columns[idx] = with_metadata |
| 1274 | |
| 1275 | if len(modified_columns) > 0: |
| 1276 | columns = [] |
| 1277 | fields = [] |
| 1278 | for i in range(len(table.schema)): |
| 1279 | if i in modified_columns: |
| 1280 | columns.append(modified_columns[i]) |
| 1281 | fields.append(modified_fields[i]) |
| 1282 | else: |
| 1283 | columns.append(table[i]) |
| 1284 | fields.append(table.schema[i]) |
| 1285 | return pa.Table.from_arrays(columns, schema=pa.schema(fields)) |
| 1286 | else: |