@Override
  public DataSource getDataSource(String column) {
    IndexContainer indexContainer = _indexContainerMap.get(column);
    if (indexContainer != null) {
      // Physical column
      return indexContainer.toDataSource();
    } else {
      // Virtual column
      FieldSpec fieldSpec = _schema.getFieldSpecFor(column);
      Preconditions.checkState(fieldSpec != null && fieldSpec.isVirtualColumn(), "Failed to find column: %s", column);
      // TODO: Refactor virtual column provider to directly generate data source
      VirtualColumnContext virtualColumnContext = new VirtualColumnContext(fieldSpec, _numDocsIndexed);
      VirtualColumnProvider virtualColumnProvider = VirtualColumnProviderFactory.buildProvider(virtualColumnContext);
      return new ImmutableDataSource(virtualColumnProvider.buildMetadata(virtualColumnContext),
          virtualColumnProvider.buildColumnIndexContainer(virtualColumnContext));
    }
  }
@Test
  public void testDataSourceForSVColumns()
      throws IOException {
    for (FieldSpec fieldSpec : _schema.getAllFieldSpecs()) {
      if (fieldSpec.isSingleValueField()) {
        String column = fieldSpec.getName();
        DataSource actualDataSource = _mutableSegmentImpl.getDataSource(column);
        DataSource expectedDataSource = _immutableSegment.getDataSource(column);

        int actualNumDocs = actualDataSource.getDataSourceMetadata().getNumDocs();
        int expectedNumDocs = expectedDataSource.getDataSourceMetadata().getNumDocs();
        assertEquals(actualNumDocs, expectedNumDocs);

        Dictionary actualDictionary = actualDataSource.getDictionary();
        Dictionary expectedDictionary = expectedDataSource.getDictionary();
        assertEquals(actualDictionary.length(), expectedDictionary.length());

        // Allow the segment name to be different
        if (column.equals(CommonConstants.Segment.BuiltInVirtualColumn.SEGMENTNAME)) {
          continue;
        }

        ForwardIndexReader actualReader = actualDataSource.getForwardIndex();
        ForwardIndexReader expectedReader = expectedDataSource.getForwardIndex();
        try (ForwardIndexReaderContext actualReaderContext = actualReader.createContext();
            ForwardIndexReaderContext expectedReaderContext = expectedReader.createContext()) {
          for (int docId = 0; docId < expectedNumDocs; docId++) {
            int actualDictId = actualReader.getDictId(docId, actualReaderContext);
            int expectedDictId = expectedReader.getDictId(docId, expectedReaderContext);
            assertEquals(actualDictionary.get(actualDictId), expectedDictionary.get(expectedDictId));
          }
        }
      }
    }
  }