in nl2sql/datasets/base.py [0:0]
def prune_schema(self) -> BaseDatasetSchema:
"""
Reduces the schema to only contain the keys present in the provided IDs
"""
schema: BaseDatasetSchema = {}
for curr_id in self.ids:
dbname, tabname, colname = curr_id.split(".")
if dbname not in schema:
schema[dbname] = {}
if tabname not in schema[dbname]:
schema[dbname][tabname] = {}
if colname not in schema[dbname][tabname]:
schema[dbname][tabname][colname] = self.dataset_schema[dbname][tabname][
colname
]
return schema