Skip to content

Commit 1f64182

Browse files
rustyconoverclaude
andcommitted
Add debug logging for batch metadata in table-in-out functions
Log input and output metadata in _process_batches to help debug FINALIZE and NEED_MORE_INPUT protocol states: - batch_received now includes input_metadata from client - batch_written now includes output_metadata being sent Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
1 parent 532c462 commit 1f64182

1 file changed

Lines changed: 26 additions & 1 deletion

File tree

vgi/worker.py

Lines changed: 26 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -852,10 +852,23 @@ def _process_batches(
852852

853853
batch_count += 1
854854
total_input_rows += batch.num_rows
855+
856+
# Log input metadata for protocol debugging
857+
input_metadata_dict = (
858+
{
859+
k.decode() if isinstance(k, bytes) else k: (
860+
v.decode() if isinstance(v, bytes) else v
861+
)
862+
for k, v in metadata.items()
863+
}
864+
if metadata
865+
else None
866+
)
855867
fn_log.debug(
856868
"batch_received",
857869
batch_index=batch_count,
858870
input_rows=batch.num_rows,
871+
input_metadata=input_metadata_dict,
859872
)
860873

861874
output = generator.send(
@@ -866,13 +879,25 @@ def _process_batches(
866879
assert output.batch is not None
867880
output_rows = output.batch.num_rows
868881
total_output_rows += output_rows
882+
883+
# Get output metadata for logging
884+
output_custom_metadata = output.metadata(invocation)
869885
writer.write_batch(
870-
output.batch, custom_metadata=output.metadata(invocation)
886+
output.batch, custom_metadata=output_custom_metadata
871887
)
888+
889+
# Log output metadata for protocol debugging
890+
output_metadata_dict = {
891+
k.decode() if isinstance(k, bytes) else k: (
892+
v.decode() if isinstance(v, bytes) else v
893+
)
894+
for k, v in output_custom_metadata.items()
895+
}
872896
fn_log.debug(
873897
"batch_written",
874898
batch_index=batch_count,
875899
output_rows=output_rows,
900+
output_metadata=output_metadata_dict,
876901
)
877902
except (KeyboardInterrupt, SystemExit):
878903
raise # Let these propagate normally

0 commit comments

Comments
 (0)