Memory leak with emit dataframe

[tomuben]

Steps to reproduce

1. Create UDF

--/
CREATE OR REPLACE PYTHON3 SET SCRIPT generate_data(
  "id" INTEGER, 
  "batch_count" INTEGER, 
  "batch_size" INTEGER
) EMITS (
  SAMPLE01 DOUBLE, 
  SAMPLE02 DOUBLE, 
  SAMPLE03 DOUBLE, 
  SAMPLE04 DOUBLE, 
  SAMPLE05 DOUBLE, 
  SAMPLE06 DOUBLE, 
  SAMPLE07 DOUBLE, 
  SAMPLE08 DOUBLE, 
  SAMPLE09 DOUBLE, 
  SAMPLE10 DOUBLE,
  SAMPLE11 DOUBLE, 
  SAMPLE12 DOUBLE, 
  SAMPLE13 DOUBLE, 
  SAMPLE14 DOUBLE, 
  SAMPLE15 DOUBLE, 
  SAMPLE16 DOUBLE, 
  SAMPLE17 DOUBLE, 
  SAMPLE18 DOUBLE, 
  SAMPLE19 DOUBLE, 
  SAMPLE20 DOUBLE,
  TARGET01 DOUBLE, 
  TARGET02 DOUBLE, 
  TARGET03 DOUBLE, 
  TARGET04 DOUBLE, 
  TARGET05 DOUBLE
) AS
import numpy as np
import pandas as pd
n_features = 20
n_targets = 5
def run(ctx):

  for i in range(ctx.batch_count):
        features = np.ones([ctx.batch_size, n_features+n_targets])
        df_result = pd.DataFrame(features)
        ctx.emit(df_result)
/

2. Run UDF:

SELECT generate_data_oom("id", 100000, 1000)
FROM VALUES BETWEEN 1 AND 2 "id"("id")
GROUP BY "id"

=> Throws an OOM error:

024-05-20	FAILED	SELECT	102.246	(null)	0	[Code: 0, SQL State: 40020]  Connection lost after system running out of memory. SessionID:1799586813378101248	

Hints:

Try https://docs.python.org/3/library/gc.html#gc.set_debug or dump objects with https://docs.python.org/3/library/gc.html#gc.get_objects

Possible root cause: Missing decrement of reference counter of a Python object.