def select_data()

in sagemaker/source/dl_utils/dataset.py [0:0]

• 27 lines of code
• 2 McCabe index (conditional complexity)

    def select_data(self, **kwargs):
        '''
        Function to select data.
        
        Parameters:
        -----------
        **kwargs: {}
            Dictionary where keys are the column names and values are the target selection.
        
        Returns:
        --------
        named_tuple with data, labels, and vehicle_properties.
        
        Usage:
        ------
        To select freightliners,
        
        > train_ds.select_data(make="FREIGHTLINER")
        
        To select 2016 freightliners that use diesel
        
        > train_ds.select_data(make="FREIGHTLINER", model_year=2016, fuel_type="diesel")
        '''
        selection_indexes = np.ones(shape=(self.vehicle_properties.shape[0]))
        for header, selection in kwargs.items():
            selection_index = self.vehicle_properties[header] == selection
            selection_indexes = np.logical_and(selection_indexes, selection_index)
        
        selected_data = self.data[selection_indexes]
        selected_labels = self.labels[selection_indexes]
        selected_vehicle_properties = self.vehicle_properties[selection_indexes]
        Selected_data = collections.namedtuple('Selected_data', 'data labels vehicle_properties')
        return Selected_data(data=selected_data, labels=selected_labels, 
                             vehicle_properties=selected_vehicle_properties)