Update app.py

app.py

@@ -1225,12 +1225,17 @@ def handle_large_dataset(df, create_document,isDataFrame):
         if len(df.columns) > 1:
                 # Skipping the original first column
                 # List of required columns
+                
                 required_columns = ['BROKER', 'ACCOUNT NUMBER', 'EMPLOYEE NAME', 'ACCOUNT NAME', 'ACCOUNT ID']
                 # Filter the DataFrame to include only the required columns
                 #print(df[required_columns])
                 #limited_data = df[required_columns]
-                limited_data11 = df.head(3)
-                limited_data = limited_data11[required_columns]
+                limited_data_top3 = df.head(3)
+                if len(df.columns) > 5:
+                    limited_data = limited_data_top3[required_columns]
+                else:
+                    limited_data = limited_data_top3
+                
                 limited_data_without_first_column = limited_data.iloc[:, 1:] 
         else:
                 limited_data = df.head(20)