# This file is part of Sympathy for Data.
# Copyright (c) 2022, Combine Control Systems AB
#
# SYMPATHY FOR DATA COMMERCIAL LICENSE
# You should have received a link to the License with Sympathy for Data.
from sylib.table_importer_gui import TableImportWidget, TableImportController
from sylib.table_sources import TableSourceModel, PreviewWorker
from sympathy.api import node as synode
from sympathy.api import importers
from sympathy.api import table
from sympathy.api import qt2 as qt_compat
from sympathy.api.exceptions import SyDataError
QtGui = qt_compat.import_module('QtGui')
QtCore = qt_compat.QtCore
QtWidgets = qt_compat.import_module('QtWidgets')
def _parquet():
from sylib.parquet import parquet
return parquet
class ImporterParquet:
"""Importer class for data in parquet format."""
def __init__(self, pq_source):
self._source = pq_source
def import_parquet(self, out_table, nr_data_rows, columns):
self._discard = False
data_table = self._source.read(nr_data_rows, columns=columns)
out_table.update(data_table)
class TableSourceParquet:
"""
This class is the layer between the physical parquet file and the import
routines.
"""
def __init__(self, fq_infilename):
self._fq_infilename = fq_infilename
def read(self, no_rows, columns=None):
out_table = table.File()
pq_table = _parquet().read_pqfile_to_table(
self._fq_infilename, rows=no_rows, columns=columns
)
names = pq_table.column_names()
for column in names:
data = pq_table.get_column_to_array(column)
out_table.set_column_from_array(column, data)
out_table.set_attributes(pq_table.get_attributes())
out_table.set_name(pq_table.get_name())
return out_table
def get_columns(self):
columns = _parquet().get_pqfile_columns(self._fq_infilename)
return columns
class TableSourceModelParquet(TableSourceModel):
"""Model layer between GUI and parquet importer."""
get_preview = qt_compat.Signal(int, list)
def __init__(self, parameters, fq_infilename, mode, valid):
super().__init__(parameters, fq_infilename, mode)
self.data_table = None
self._valid = valid
self._pq_source = TableSourceParquet(fq_infilename)
self._importer = ImporterParquet(self._pq_source)
self._init_model_specific_parameters()
self._init_preview_worker()
def _init_model_common_parameters(self):
"""Init common parameters xlsx, csv, mat and parquet importers."""
self.preview_start_row = self._parameters['preview_start_row']
self.no_preview_rows = self._parameters['no_preview_rows']
def _init_model_specific_parameters(self):
"""Init special parameters for parquet importer."""
self.preview_start_row = self._parameters['preview_start_row']
self.no_preview_rows = self._parameters['no_preview_rows']
self.data_offset = lambda: None
self.data_offset.value = 1
self.columns = self._parameters["columns"]
self.nr_rows = self._parameters["nr_rows"]
if self._valid:
self._cnames = self._pq_source.get_columns()
self.columns.list = list(self._cnames)
def _init_preview_worker(self):
self._preview_thread = QtCore.QThread()
self._preview_worker = PreviewWorker(self._importer.import_parquet)
self._preview_worker.moveToThread(self._preview_thread)
self._preview_thread.finished.connect(self._preview_worker.deleteLater)
self.get_preview.connect(self._preview_worker.create_preview_table)
self._preview_worker.preview_ready.connect(self.set_preview_table)
self._preview_worker.preview_failed.connect(self.set_preview_failed)
self._preview_thread.start()
self.collect_preview_values()
@qt_compat.Slot()
def collect_preview_values(self):
"""Collect preview data from parquet file."""
no_rows = self.no_preview_rows.value
if self._valid:
# If another file with different columns was previously configured,
# it would leave behind a different selection. Trying to match it
# against the new columns then causes the exception handled below.
try:
columns = self.columns.selected_names(self._cnames)
except SyDataError:
columns = [col for col in self._cnames
if col in self.columns.value_names]
self.get_preview.emit(no_rows, columns)
else:
self.data_table = table.File()
self.update_table.emit()
@qt_compat.Slot(table.File)
def set_preview_table(self, data_table):
self.data_table = data_table
self.update_table.emit()
def cleanup(self):
self._preview_thread.quit()
self._preview_thread.wait()
class ImportParametersWidgetParquet(QtWidgets.QWidget):
"""
The control group box includes the widgets for determination of
data start and end row/column and transpose condition.
"""
get_preview = QtCore.Signal()
def __init__(self, model, parent=None):
super().__init__(parent)
self._model = model
self._init_gui(model)
self._init_preview_signals()
def _init_gui(self, model):
self._columns = model.columns.gui()
layout = QtWidgets.QVBoxLayout()
layout.setAlignment(QtCore.Qt.AlignmentFlag.AlignLeft)
layout.addWidget(model.nr_rows.gui())
layout.addWidget(model.columns.gui())
self._update_button = QtWidgets.QPushButton("Update Preview")
layout.addWidget(self._update_button)
self.setLayout(layout)
def _init_preview_signals(self):
self._update_button.clicked.connect(self.get_preview)
class TableImportWidgetParquet(TableImportWidget):
MODE = 'Parquet'
def __init__(self, parameters, fq_infilename, valid=True):
self.model = TableSourceModelParquet(
parameters, fq_infilename, self.MODE, valid
)
super().__init__(parameters, fq_infilename, self.MODE, valid)
def _collect_import_parameters_widget(self, model):
return ImportParametersWidgetParquet(model)
def _collect_table_source_widget(self, model):
pass
def _collect_controller(self, **kwargs):
return TableImportController(**kwargs)
[docs]
class DataImportParquet(importers.TableDataImporterBase):
"""Importer for Parquet files."""
IMPORTER_NAME = "Parquet"
def __init__(self, fq_infilename, parameters):
super().__init__(fq_infilename, parameters)
if parameters is not None:
self._init_parameters()
def _init_parameters(self):
parameters = self._parameters
if "preview_start_row" not in parameters:
parameters.set_integer(
"preview_start_row",
value=1,
label="Preview start row",
description="The first row where data will review from.",
editor=synode.editors.bounded_spinbox_editor(1, 500, 1),
)
if "no_preview_rows" not in parameters:
parameters.set_integer(
"no_preview_rows",
value=20,
label="Number of preview rows",
description="The number of preview rows to show.",
editor=synode.editors.bounded_spinbox_editor(1, 200, 1),
)
if "nr_rows" not in parameters:
parameters.set_integer(
"nr_rows",
value=0,
label="Number of rows to be imported",
description=(
"How many rows are to be imported from the parquet file. "
"Entering 0 (the default) will import all rows. "
),
editor=synode.editors.bounded_spinbox_editor(0, None, 1),
)
if "columns" not in parameters:
editor = synode.editors.multilist_editor(edit=True)
parameters.set_list(
"columns",
label="Select columns to be imported",
description=(
"Only the selected columns (if present in the file) "
"will be imported."
),
value=[],
editor=editor,
)
parameters["columns"].multiselect_mode = "passthrough"
def name(self):
return self.IMPORTER_NAME
def valid_for_file(self):
"""Return True if input file is a valid Parquet file."""
return _parquet().is_parquet_file(self._fq_infilename)
def parameter_view(self, parameters):
valid_for_file = self.valid_for_file()
return TableImportWidgetParquet(
parameters, self._fq_infilename, valid_for_file
)
def import_data(self, out_datafile, parameters=None, progress=None):
"""Import Parquet data from a file"""
parameters = parameters
# Establish connection to Parquet datasource
table_source = TableSourceParquet(self._fq_infilename)
column_names = table_source.get_columns()
columns = parameters["columns"].selected_names(column_names)
nr_rows = parameters["nr_rows"].value
importer = ImporterParquet(table_source)
try:
importer.import_parquet(out_datafile, nr_rows, columns)
except Exception as e:
raise self.import_failed(e) from e
