def test_cast_dataframe_types_no_fields(self):
mock_fields = []
mock_known_representations = Mock(name='known_representations')
schema = RecordsSchema(
fields=mock_fields,
known_representations=mock_known_representations)
mock_df = Mock(name='df')
out = schema.cast_dataframe_types(mock_df)
self.assertEqual(out, mock_df.apply.return_value)
def test_to_schema_sql(self, mock_schema_to_schema_sql):
mock_driver = Mock(name='driver')
mock_schema_name = Mock(name='schema_name')
mock_table_name = Mock(name='table_name')
obj = RecordsSchema(fields=[], known_representations={})
out = obj.to_schema_sql(mock_driver, mock_schema_name, mock_table_name)
mock_schema_to_schema_sql.assert_called_with(
driver=mock_driver,
records_schema=obj,
schema_name=mock_schema_name,
table_name=mock_table_name)
self.assertEqual(out, mock_schema_to_schema_sql.return_value)
def test_assign_dataframe_names_no_index(self):
data = [{'a': 1}]
df = DataFrame.from_dict(data)
mock_field_a = Mock(name='field_a')
mock_field_a.name = 'mya'
mock_fields = [mock_field_a]
mock_known_representations = Mock(name='known_representations')
schema = RecordsSchema(
fields=mock_fields,
known_representations=mock_known_representations)
out = schema.assign_dataframe_names(False, df)
self.assertEqual(out.to_dict(orient='records'), [{'mya': 1}])
def test_assign_dataframe_names_with_index(self):
data = [{'b': 1}]
df = DataFrame.from_dict(data)
self.assertEqual(df.to_dict(orient='index'), {0: {'b': 1}})
mock_field_a = Mock(name='field_a')
mock_field_a.name = 'mya'
mock_field_b = Mock(name='field_b')
mock_field_b.name = 'myb'
mock_fields = [mock_field_a, mock_field_b]
mock_known_representations = Mock(name='known_representations')
schema = RecordsSchema(
fields=mock_fields,
known_representations=mock_known_representations)
out = schema.assign_dataframe_names(True, df)
self.assertEqual(out.to_dict(orient='records'), [{'myb': 1}])
self.assertEqual(out.to_dict(orient='index'), {'mya': {'myb': 1}})
def test_refine_from_dataframe(self, mock_refine_schema_from_dataframe):
mock_fields = Mock(name='fields')
mock_known_representations = Mock(name='known_representations')
schema = RecordsSchema(
fields=mock_fields,
known_representations=mock_known_representations)
mock_df = Mock(name='df')
mock_processing_instructions = Mock(name='processing_instructions')
out = schema.refine_from_dataframe(mock_df,
mock_processing_instructions)
mock_refine_schema_from_dataframe.\
assert_called_with(records_schema=schema,
df=mock_df,
processing_instructions=mock_processing_instructions)
self.assertEqual(out, mock_refine_schema_from_dataframe.return_value)
def schema_from_dataframe(df: DataFrame,
processing_instructions: ProcessingInstructions,
include_index: bool) -> 'RecordsSchema':
from records_mover.records.schema import RecordsSchema # noqa
from records_mover.records.schema.field import RecordsSchemaField # noqa
fields = []
origin_representation = \
RecordsSchemaKnownRepresentation.from_dataframe(df, processing_instructions)
known_representations: Dict[str, RecordsSchemaKnownRepresentation] = {
'origin': origin_representation
}
if include_index:
fields.append(
RecordsSchemaField.from_index(
df.index, processing_instructions=processing_instructions))
for column in df:
fields.append(
RecordsSchemaField.from_series(
df[column], processing_instructions=processing_instructions))
return RecordsSchema(fields=fields,
known_representations=known_representations)
def refine_schema_from_dataframe(records_schema: 'RecordsSchema',
df: DataFrame,
processing_instructions:
ProcessingInstructions = ProcessingInstructions()) ->\
'RecordsSchema':
from records_mover.records.schema import RecordsSchema
max_sample_size = processing_instructions.max_inference_rows
total_rows = len(df.index)
if max_sample_size is not None and max_sample_size < total_rows:
sampled_df = df.sample(n=max_sample_size)
else:
sampled_df = df
rows_sampled = len(sampled_df.index)
fields = [
field.refine_from_series(sampled_df[field.name],
total_rows=total_rows,
rows_sampled=rows_sampled)
for field in records_schema.fields
]
return RecordsSchema(
fields=fields,
known_representations=records_schema.known_representations)
def test_str(self):
obj = RecordsSchema(fields=[], known_representations={})
self.assertEqual(str(obj), "RecordsSchema(types={})")