-- jsonb decoding/transformation
CREATE OR REPLACE FUNCTION py_jsonb_length_json(val_str jsonb) RETURNS int
LANGUAGE plpythonu
AS $$
  import json
  val = json.loads(val_str)
  return len(val)
$$;

CREATE OR REPLACE FUNCTION py_jsonb_length_trans(val jsonb) RETURNS int
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
  return len(val)
$$;

-- returning numeric arrays
CREATE OR REPLACE FUNCTION py_jsonb_ret_int_array_trans() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
  return [0,1,2,3,4,5,6,7,8,9]
$$;

CREATE OR REPLACE FUNCTION py_jsonb_ret_float_array_trans() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
  return [0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0]
$$;

CREATE OR REPLACE FUNCTION py_jsonb_ret_int_array_json() RETURNS jsonb
LANGUAGE plpythonu
AS $$
  import json
  return json.dumps([0,1,2,3,4,5,6,7,8,9])
$$;

CREATE OR REPLACE FUNCTION py_jsonb_ret_int_array_raw() RETURNS jsonb
LANGUAGE plpythonu
AS $$
  return [0,1,2,3,4,5,6,7,8,9]
$$;

CREATE OR REPLACE FUNCTION py_jsonb_ret_float_array_json() RETURNS jsonb
LANGUAGE plpythonu
AS $$
  import json
  return json.dumps([0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0])
$$;

CREATE OR REPLACE FUNCTION py_jsonb_ret_float_array_raw() RETURNS jsonb
LANGUAGE plpythonu
AS $$
  return [0.0,1.0,2.0,3.0,4.0,5.0,6.0,7.0,8.0,9.0]
$$;

-- returning objects
CREATE OR REPLACE FUNCTION py_jsonb_ret_obj_trans() RETURNS jsonb
LANGUAGE plpythonu
AS $$
  import json
  return json.dumps({"aaa": 123, "bbb": 456})
$$;


CREATE OR REPLACE FUNCTION py_jsonb_ret_obj_trans() RETURNS jsonb
LANGUAGE plpythonu
TRANSFORM FOR TYPE jsonb
AS $$
  return {"aaa": 123, "bbb": 456}
$$;


-- Tests
DO $$
BEGIN
  FOR i IN 1..1000000 LOOP
    PERFORM py_jsonb_length_trans('{"aaa": 123, "bbb": 456}'::jsonb);
  END LOOP;
END
$$;

DO $$
BEGIN
  FOR i IN 1..1000000 LOOP
    PERFORM py_jsonb_length_json('{"aaa": 123, "bbb": 456}'::jsonb);
  END LOOP;
END
$$;

DO $$
BEGIN
  FOR i IN 1..1000000 LOOP
    PERFORM py_jsonb_ret_int_array_trans();
  END LOOP;
END
$$;

DO $$
BEGIN
  FOR i IN 1..1000000 LOOP
    PERFORM py_jsonb_ret_int_array_raw();
  END LOOP;
END
$$;

DO $$
BEGIN
  FOR i IN 1..1000000 LOOP
    PERFORM py_jsonb_ret_int_array_json();
  END LOOP;
END
$$;
