You might be able to do
prepare build(select: ([arg(column_name)])
otherwise this should work
prepare fn query, context ->
column_name = Ash.Query.get_argument(query, :column_name)
query
|> Ash.Query.select([column_name])
end