Spaces:
Runtime error
Runtime error
File size: 1,717 Bytes
e4f9cbe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 |
"""Tests for the pandas source."""
import pandas as pd
from ...schema import schema
from .pandas_source import PANDAS_INDEX_COLUMN, PandasDataset
from .source import SourceSchema
def test_simple_dataframe() -> None:
df = pd.DataFrame.from_records([{
'name': 'a',
'age': 1
}, {
'name': 'b',
'age': 2
}, {
'name': 'c',
'age': 3
}])
source = PandasDataset(df)
source.setup()
source_schema = source.source_schema()
assert source_schema == SourceSchema(
fields=schema({
PANDAS_INDEX_COLUMN: 'int64',
'name': 'string',
'age': 'int64'
}).fields,
num_items=3)
items = list(source.process())
assert items == [{
PANDAS_INDEX_COLUMN: 0,
'name': 'a',
'age': 1
}, {
PANDAS_INDEX_COLUMN: 1,
'name': 'b',
'age': 2
}, {
PANDAS_INDEX_COLUMN: 2,
'name': 'c',
'age': 3
}]
def test_simple_dataframe_with_index() -> None:
df = pd.DataFrame.from_records([{
'name': 'a',
'age': 1
}, {
'name': 'b',
'age': 2
}, {
'name': 'c',
'age': 3
}],
index=['id1', 'id2', 'id3'])
source = PandasDataset(df)
source.setup()
source_schema = source.source_schema()
assert source_schema == SourceSchema(
fields=schema({
PANDAS_INDEX_COLUMN: 'string',
'name': 'string',
'age': 'int64'
}).fields,
num_items=3)
items = list(source.process())
# The PANDAS_INDEX_COLUMN aligns with the pandas index.
assert items == [{
PANDAS_INDEX_COLUMN: 'id1',
'name': 'a',
'age': 1
}, {
PANDAS_INDEX_COLUMN: 'id2',
'name': 'b',
'age': 2
}, {
PANDAS_INDEX_COLUMN: 'id3',
'name': 'c',
'age': 3
}]
|