Вызов Snowflake UDTF из сохраненной процедуры снежинки ⇐ Python

[Anonymous]

Я пытаюсь написать UDTF в Snowflake, который мне нужно позвонить из хранимой процедуры. Я хотел бы зарегистрировать UDTF как часть сеанса, используя SnowPark. Кроме того, хочу назвать его с разделением, поэтому он работает для каждого раздела вместо того, чтобы работать для каждой записи, потому что я хочу запустить трансформацию, последовательно внутри каждого раздела, чтобы контролировать уровень параллелизма. 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): < /p>
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"
}
}
< /code>
Ниже приведен код обработчика для этой хранимой процедуры, которая вызывает функцию таблицы: < /p>
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': '/encrypted_test/sample_plaintext.txt.pgp'
},
{
'FILE_URL': '/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()
< /code>
И, наконец, функция таблицы: < /p>
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
< /code>
Ниже приведена некоторая информация из документации Snowflake, в которой упоминается о наличии либо end_partition, либо процесса. Так что не уверен, что пойдет не так
https://docs.snowflake.com/en/developer ... tabularize>

Подробнее здесь: https://stackoverflow.com/questions/796 ... -procedure