370def extract_buffers(events, classified):
371 """Extract flat numpy buffers from an awkward array.
372
373 Returns a dict mapping column name -> numpy array, covering all container
374 offsets, nested-vector offsets, and input data columns. Output column
375 buffers are not included (allocate_outputs handles those).
376
377 Parameters
378 ----------
379 events:
380 An ak.Array (real or zero-length after typetracer conversion).
381 classified:
382 Output of classify_columns or resolve_optional_columns.
383 """
384 buffers = {}
385 num_events = int(ak.num(events, axis=0))
386 synthesized_offsets =
set()
387
388
389
390 deferred_links = []
391
392 for container_name, info in classified.items():
393 nested_offsets = info["nested_offsets"]
394
395
396
397 for nested_offset_name, nested in nested_offsets.items():
398 for col in nested["inputs"]:
399 if col.is_variant_link:
400 raise NotImplementedError(
401 f"variant link columns are not supported "
402 f"(column '{col.name}')"
403 )
404 if _is_link_column(col):
405 deferred_links.append((nested_offset_name, col))
406 continue
407 base = _branch_name_for_column(col.name)
408 inner = _inner_most_list_offset_array(events[base])
409 raw_offsets = np.asarray(inner.layout.offsets.data)
410 start = int(raw_offsets[0])
411 end = int(raw_offsets[-1])
412
413
414 buffers[nested_offset_name] = np.ascontiguousarray(
415 raw_offsets - start, dtype=np.uint64
416 )
417 buffers[col.name] = np.ascontiguousarray(
418 inner.layout.content.data[start:end]
419 )
420
421 for col in info["inputs"]:
422 if col.is_variant_link:
423 raise NotImplementedError(
424 f"variant link columns are not supported (column '{col.name}')"
425 )
426 link_inputs = [col for col in info["inputs"] if _is_link_column(col)]
427 flat_inputs = [col for col in info["inputs"] if not _is_link_column(col)]
428 deferred_links.extend((container_name, col) for col in link_inputs)
429
430 if not flat_inputs and not nested_offsets and not link_inputs:
431
432 buffers[container_name] = np.array([0, num_events], dtype=np.uint64)
433 synthesized_offsets.add(container_name)
434 continue
435
436 if not flat_inputs:
437
438
439 any_nested_input = next(
440 (col for nested in nested_offsets.values() for col in nested["inputs"]),
441 None,
442 )
443 if any_nested_input is not None:
444 base = _branch_name_for_column(any_nested_input.name)
445 jagged = events[base]
446 elif link_inputs:
447 jagged, _key = _link_index_and_key(
448 events, link_inputs[0].name, link_inputs[0].name
449 )
450 else:
451 jagged = None
452 if jagged is not None:
453 buffers[container_name] = _offsets_from_counts(
454 ak.to_numpy(ak.num(jagged, axis=1))
455 )
456 else:
457 buffers[container_name] = np.array([0, num_events], dtype=np.uint64)
458 synthesized_offsets.add(container_name)
459 continue
460
461
462
463 sorted_cols = sorted(flat_inputs, key=lambda c: c.offset_name)
464
465 for offset_name, cols_iter in itertools.groupby(
466 sorted_cols, key=lambda c: c.offset_name
467 ):
468 cols = list(cols_iter)
469 unzipped = {col.name: events[col.name] for col in cols}
470 zipped = ak.zip(unzipped)
471
472
473 form, length, raw_buffers = ak.to_buffers(
474 zipped, form_key=f"{offset_name}{{id}}"
475 )
476
477 if isinstance(form, ak.forms.RecordForm):
478
479
480
481
482 buffers[container_name] = np.array(
483 [0, length], dtype=np.uint64
484 )
485 for col in cols:
486 buffers[col.name] = ak.to_numpy(events[col.name])
487 elif isinstance(form, ak.forms.ListOffsetForm):
488
489
490 offset_key = next(
491 key for key in raw_buffers if key.endswith("-offsets")
492 )
493 buffers[container_name] = np.asarray(
494 raw_buffers[offset_key]
495 ).astype(np.uint64)
496
497
498 inner = form.content
499 for field in inner.fields:
500 buffers[field] = np.asarray(
501 raw_buffers[f"{inner.content(field).form_key}-data"]
502 )
503 else:
504 raise RuntimeError(
505 f"Cannot handle form {type(form)} for "
506 f"container {container_name}"
507 )
508
509
510 for offset_owner, col in deferred_links:
511 target = col.sole_link_target_name
512 target_offsets = buffers.get(target)
513 if target_offsets is None or target in synthesized_offsets:
514 raise RuntimeError(
515 f"link column '{col.name}' targets container '{target}', "
516 "whose offsets could not be derived from the tool's input columns"
517 )
518 if col.name.endswith(".data"):
519 nested_link_offsets, data = _convert_vector_link_column(
520 events, col, target_offsets
521 )
522 buffers[offset_owner] = nested_link_offsets
523 buffers[col.name] = data
524 else:
525 buffers[col.name] = _convert_scalar_link_column(
526 events, col, target_offsets
527 )
528
529 return buffers
530
531