[Python] Calling snowflake udtf from snowflake stored procedure

I am trying to write a udtf in snowflake which I need to call from stored procedure. I would want to register the udtf as part of session using snowpark. Also, want to call it with partition by so it runs for each partition instead of running for each record because I want to run the tranformation serially inside each partition to control the level of parallelism.

below is the code I pulled up so far, but it complains that it is not able to find process method.

Procedure error: (1304): 01bd9ed6-020e-0336-003b-dd0712547cce: 100357 (P0000): 01bd9ed6-020e-0336-003b-dd0712547cce: Could not find process method in function FILE_PROCESSOR_UDTF with handler compute"

below is my definition of stored procedure (terraform):

resource "snowflake_procedure_python" "file_processor" {
name = "FILE_PROCESSOR"
comment = "Process files using UDTF with file URL input"
database = var.database_name
schema = var.application_schema_name
snowpark_package = var.snowpark_version
runtime_version = var.python_version
handler = "transform_snowflake.procedures.file_processor.main"
# return_type = "TABLE(FILE_URL VARCHAR, FILE_SIZE NUMBER, STATUS VARCHAR, CONTENT_PREVIEW VARCHAR, ERROR_MESSAGE VARCHAR)"
return_type = "VARIANT"
execute_as = "CALLER"

arguments {
arg_name = "STAGE_PATH"
arg_data_type = "VARCHAR"
}
arguments {
arg_name = "FILE_PATTERN"
arg_data_type = "VARCHAR"
arg_default_value = "'*'"
}
arguments {
arg_name = "MAX_FILES"
arg_data_type = "NUMBER"
arg_default_value = 10
}

imports {
stage_location = var.code_stage_fully_qualified_name
path_on_stage = "transform_snowflake.zip"
}
}


below is the handler code for this stored procedure which is calling table function:

import snowflake.snowpark as snowpark
from snowflake.snowpark import Session
import pandas as pd
import logging
import traceback
from snowflake.snowpark.types import StringType, StructField, StructType, VariantType, IntegerType
from transform_snowflake.udtfs.file_processor_udtf import FileProcessorUDTF

logger = logging.getLogger("transform_snowflake.procedures.file_processor")

def main(session: Session, stage_path: str, file_pattern: str = '*', max_files: int = 10):
"""
Main stored procedure function that registers UDTF and processes files
"""
try:
logger.info(f"Starting file processor for stage: {stage_path}, pattern: {file_pattern}")

output_schema = StructType([
StructField("FILE_URL", StringType()),
StructField("FILE_SIZE", IntegerType()),
StructField("STATUS", StringType()),
StructField("CONTENT_PREVIEW", StringType()),
StructField("ERROR_MESSAGE", StringType())
])
file_processor_udtf = session.udtf.register(
handler=FileProcessorUDTF,
name="file_processor_udtf",
input_types=[StringType()],
output_schema=output_schema,
replace=True,
is_permanent=False,
packages=['snowflake-snowpark-python', 'pgpy', 'pandas']
)

# Adding two rows in dataframe for testing purposes
df_files = session.create_dataframe([{
'FILE_URL': '<stage_name>/encrypted_test/sample_plaintext.txt.pgp'
},
{
'FILE_URL': '<stage_name>/encrypted_test/sample_plaintext2.txt.pgp'
}])
# Check if we have any files
file_count = df_files.count()
if file_count == 0:
logger.warning(f"No files found in {stage_path} with pattern {file_pattern}")
# Return empty result with correct schema
return session.create_dataframe([{
'FILE_URL': 'No files found',
'FILE_SIZE': 0,
'STATUS': 'NO_FILES',
'CONTENT_PREVIEW': '',
'ERROR_MESSAGE': f'No files found in {stage_path} with pattern {file_pattern}'
}])

logger.info(f"Found {file_count} files to process")

# Use the UDTF to process files
df_columns = df_files.columns
result_df = df_files.join_table_function(
file_processor_udtf(*df_columns).over(partition_by='FILE_URL'))

# Collect and return results
return result_df.collect()

except Exception as e:
logger.error(f"Error in file processor: {str(e)}")
logger.error(traceback.format_exc())

# Return error result
return session.create_dataframe([{
'FILE_URL': stage_path,
'FILE_SIZE': 0,
'STATUS': 'PROCEDURE_ERROR',
'CONTENT_PREVIEW': '',
'ERROR_MESSAGE': f'Procedure error: {str(e)}'
}]).collect()


And lastly, the table function:

import pandas as pd
import time
import logging
from snowflake.snowpark.files import SnowflakeFile
from _snowflake import vectorized
import threading

logger = logging.getLogger("udtfs.file_processor_udtf")

class FileProcessorUDTF:
"""
UDTF class to process files from stage URLs using vectorized processing
"""

def __init__(self):
self.lock = threading.Lock()

@vectorized(input=pd.DataFrame)
def end_partition(self, df: pd.DataFrame):
"""
Vectorized end_partition function that processes files sequentially
"""
raise Exception('inside end_partition')
logger.info(f"Processing partition with {len(df)} files")

results = []

# Process files sequentially with simple for loop
for idx, row in df.iterrows():
try:
result = self._process_single_file(row['file_url'])
results.append(result)
except Exception as e:
# Handle unexpected errors in file processing
error_result = {
'file_url': row['file_url'],
'file_size': 0,
'status': 'FAILED',
'content_preview': '',
'error_message': f'Processing error: {str(e)}'
}
results.append(error_result)

# Convert results to DataFrame
result_df = pd.DataFrame(results)

logger.info(f"Successfully processed {len(results)} files in partition")
return result_df

def _process_single_file(self, file_url):
"""
Process a single file using SnowflakeFile
"""
start_time = time.time()

result = {
'FILE_URL': file_url,
'FILE_SIZE': 0,
'STATUS': 'FAILED',
'CONTENT_PREVIEW': '',
'ERROR_MESSAGE': ''
}

try:
# Open file using SnowflakeFile
with SnowflakeFile.open(file_url, 'rb') as f:
# Get file size
if f.seekable():
current_pos = f.tell()
f.seek(0, 2) # Seek to end
file_size = f.tell()
f.seek(current_pos) # Return to original position

# Read first 500 bytes for preview
preview_data = f.read(500)
f.seek(current_pos) # Reset position
else:
# If not seekable, read in chunks
preview_data = f.read(500)
file_size = len(preview_data)

# Convert preview to string (handle binary data)
try:
content_preview = preview_data.decode('utf-8', errors='replace')[:200]
except:
content_preview = f"Binary data ({len(preview_data)} bytes)"

# Update result with success
result.update({
'FILE_SIZE': file_size,
'STATUS': 'SUCCESS',
'CONTENT_PREVIEW': content_preview,
'ERROR_MESSAGE': ''
})

except Exception as e:
# Handle file processing errors
result.update({
'STATUS': 'FAILED',
'ERROR_MESSAGE': str(e)[:500] # Truncate long error messages
})

return result


Below is some information from snowflake documentation that mentions about having either end_partition or process. so not sure what is going wrong

https://docs.snowflake.com/en/developer-guide/udf/python/udf-python-tabular-vectorized

Continue reading...