Fix TSDataSampler Slicing Bug #1716

[YeewahChan]

Title: Fix TSDataSampler Slicing Bug #1716

Description

I have optimized the mechanism for reading data from data_arr in the __getitem__ method:

1. In most cases, maintain high performance of slicing and preserve the original logic.
2. For the issue raised by #1716, I adopted a straightforward approach to circumvent the problem: set self.nan_idx to len(self.data_arr) - 1 instead of -1.

Motivation and Context

Related Issue Link The __getitem__ function of TSDataSampler generally operates without issue, yet in certain scenarios, particularly when step_len is significant and nan must be added prior to start, it may inadvertently sample an empty list, which is erroneous. Consider the following code snippet as an example:

import numpy as np
import pandas as pd
from qlib.data.dataset import TSDataSampler

datetimes = [
    '2000-01-31', '2000-02-29', '2000-03-31', '2000-04-30', '2000-05-31'
]
instruments = ['000001', '000002', '000003', '000004', '000005']
index = pd.MultiIndex.from_product([pd.to_datetime(datetimes), instruments],
                                   names=['datetime', 'instrument'])
data = np.random.randn(len(datetimes) * len(instruments))
test_df = pd.DataFrame(data=data, index=index, columns=['ret'])
dataset = TSDataSampler(test_df, datetimes[0], datetimes[-1], step_len=2)
print(dataset[0])

The expected outcome should be an array with nan as the initial element followed by numerical values. However, the actual result is an empty list, as depicted in the following image:

Mechanism of the Error

How did this issue arise? It stems from the fact that the dataframe data has been stored in an ndarray (variable name: data_arr) in the order of 'instrument, time', and thus a set of indicies is generated during the fetch. The logical error occurs when step_len is large, and the count from the current processing position back to start requires nan supplementation (as shown in the following image): These nan values are then converted to -1 during subsequent processing (as shown in the following image) because self.nan_idx is hardcoded as -1: Therefore, the error occurs when there are several -1s at the beginning, causing the slice to become [-1, index of the last value to be fetched]. In cases requiring padding, this index is typically small, leading to the retrieval of an empty list, which is the error reported in the issue.

My Modification Approach

Since the error always occurs at the beginning with nan, the indicies[0] must be self.nan_idx, which is hardcoded to -1. For this scenario, it causes issues with the retrieved slices. After examining, I found that self.nan_idx is only called once in this context. Therefore, following the original author's intention, I changed self.nan_idx to len(self.data_arr) - 1. This solution not only addresses the problem but also avoids reducing code readability.

How Has This Been Tested?

• [x] Passed the test by running: pytest qlib/tests/test_all_pipeline.py under the upper directory of qlib.
• [x] If you are adding a new feature, tested on your own test scripts.

My test scripts are located in tests/data_mid_layer_tests/test_dataset.py. TestTSDataSampler includes two test contents:

1. Reproducing and testing the resolution of #1716 (i.e., partial sampling results in an empty list).
2. An additional test case should be considered: if the first element does not require nan padding and is directly fetched, the test result should be the direct retrieval of the first value.

Screenshots of Test Results (if appropriate):

1. Pipeline test:

2. Your own tests:

Types of changes

• [x] Bug fix
• [ ] New feature
• [ ] Documentation update

[YeewahChan]

@microsoft-github-policy-service agree

[you-n-g]

Could you please fix the CI errors? @YeewahChan

[YeewahChan]

Could you please fix the CI errors? @YeewahChan

The CI errors were due to pylint formatting issues, but I've fixed them now.

[you-n-g]

Thanks for your great efforts!