Coverage for mindsdb / api / executor / sql_query / steps / sql_steps.py: 40%

1import pandas as pd 

2 

3from mindsdb.api.executor.planner.steps import ( 

4 LimitOffsetStep, 

5 DataStep, 

6) 

7 

8from mindsdb.api.executor.sql_query.result_set import ResultSet 

9 

10from .base import BaseStepCall 

11 

12 

13class LimitOffsetStepCall(BaseStepCall): 

14 

15 bind = LimitOffsetStep 

16 

17 def call(self, step): 

18 step_data = self.steps_data[step.dataframe.step_num] 

19 

20 df = step_data.get_raw_df() 

21 

22 step_data2 = ResultSet(columns=list(step_data.columns)) 

23 

24 if isinstance(step.offset, int): 

25 df = df[step.offset:] 

26 if isinstance(step.limit, int): 

27 df = df[:step.limit] 

28 

29 step_data2.add_raw_df(df) 

30 

31 return step_data2 

32 

33 

34class DataStepCall(BaseStepCall): 

35 

36 bind = DataStep 

37 

38 def call(self, step): 

39 # create resultset 

40 df = pd.DataFrame(step.data) 

41 return ResultSet.from_df(df, database='', table_name='')