get_data_matrix function

Helper function to extract data matrix with proper handling