Fixed start_col variable (#35)

Author: umitbuyuksahin

exasol_udf_mock_python/mock_context.py

@@ -61,12 +61,13 @@ def get_dataframe(self, num_rows='all', start_col=0):
             raise RuntimeError("get_dataframe() parameter 'start_col' must be an integer >= 0")
         if self._data is None:
             return None
-        columns_ = [column.name for column in self._metadata.input_columns]
+        columns_ = [column.name for column in self._metadata.input_columns[start_col:]]
 
         i = 0
         df = None
         while num_rows == 'all' or i < num_rows:
-            df_current = pd.DataFrame.from_records([self._data], columns=columns_)
+            df_current = pd.DataFrame.from_records(
+                [self._data[start_col:]], columns=columns_)
             if df is None:
                 df = df_current
             else:

tests/test_executor_context_set_emits.py

@@ -311,6 +311,7 @@ def run(ctx):
     with pytest.raises(RuntimeError) as excinfo:
         result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
 
+
 def test_get_dataframe_start_col_0():
     def udf_wrapper():
         def run(ctx):
@@ -329,23 +330,27 @@ def run(ctx):
     result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
     assert result == [Group([(1,), ])]
 
+
 def test_get_dataframe_start_col_positive():
     def udf_wrapper():
         def run(ctx):
             df = ctx.get_dataframe(num_rows=1, start_col=1)
             ctx.emit(df)
 
-    executor = UDFMockExecutor()
     meta = MockMetaData(
         script_code_wrapper_function=udf_wrapper,
         input_type="SET",
-        input_columns=[Column("t", int, "INTEGER")],
+        input_columns=[
+            Column("t1", int, "INTEGER"),
+            Column("t2", int, "INTEGER")],
         output_type="EMITS",
         output_columns=[Column("t", int, "INTEGER")]
     )
+    executor = UDFMockExecutor()
     exa = MockExaEnvironment(meta)
-    result = executor.run([Group([(1,), (2,), (3,), (4,), (5,), (6,)])], exa)
-    assert result == [Group([(1,), ])]
+    result = executor.run([Group([(1, -1), (2, -2), (3, -3), (4, -4)])], exa)
+    assert result == [Group([(-1,), ])]
+
 
 def test_emit_tuple_exception():
     def udf_wrapper():