code: def test_query_data(self):
      # single step
      self.generate_data()
      df = self._discrete_read_data()
      df = df.toPandas()
      self.verify_discrete_single_step_except_rewards(df)
      self.assertEq(
          df["reward"], np.array([0.0, 1.0, 4.0, 5.0], dtype="float32")
      )
      logger.info("discrete single-step seems fine")
  
      # single step with reward := reward^3 + 10
      df = self._discrete_read_data(
          custom_reward_expression="POWER(reward, 3) + 10"
      )
      df = df.toPandas()
      self.verify_discrete_single_step_except_rewards(df)
      self.assertEq(
          df["reward"],
          np.array([10.0, 11.0, 74.0, 135.0], dtype="float32"),
      )
      logger.info("discrete single-step custom reward seems fine")
  
      # multi-step
      gamma = 0.9
      self.generate_data(multi_steps=True)
      df = self._discrete_read_data(multi_steps=2, gamma=gamma)
      df = df.toPandas()
      self.verify_discrete_multi_steps_except_rewards(df)
      self.assertAllClose(
          df["reward"],
          np.array(
              [
                  gamma * 1,
                  1 * 1.0 + gamma * 4,
                  1 * 4.0 + gamma * 5,
                  1 * 5.0,
              ],
              dtype="float32",
          ),
      )
      logger.info("discrete multi-step seems fine.")
  
file_path: ReAgent/reagent/test/workflow/test_query_data.py
imports: import logging
  import unittest
  import numpy
  import pytest
  from pyspark.sql.functions import asc
  from reagent.data.oss_data_fetcher import OssDataFetcher
  from reagent.test.test_data.ex_mdps import generate_discrete_mdp_pandas_df
  from reagent.test.workflow.reagent_sql_test_base import ReagentSQLTestBase
  from reagent.workflow.types import Dataset, TableSpec
label_desc: Use image node_id 4 for calling the TestQueryData obj's underlying member method code with example usage: obj.test_query_data() without return types