mirror of
https://github.com/apache/superset.git
synced 2026-04-14 05:34:38 +00:00
feat: new Columnar upload form and API (#28192)
This commit is contained in:
Daniel Vaz Gaspar
committed by
GitHub
GitHub
parent
f5843fe588
commit
9a339f08a7
253
tests/unit_tests/commands/databases/columnar_reader_test.py
Normal file
253
tests/unit_tests/commands/databases/columnar_reader_test.py
Normal file
@@ -0,0 +1,253 @@
|
||||
# Licensed to the Apache Software Foundation (ASF) under one
|
||||
# or more contributor license agreements. See the NOTICE file
|
||||
# distributed with this work for additional information
|
||||
# regarding copyright ownership. The ASF licenses this file
|
||||
# to you under the Apache License, Version 2.0 (the
|
||||
# "License"); you may not use this file except in compliance
|
||||
# with the License. You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing,
|
||||
# software distributed under the License is distributed on an
|
||||
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
||||
# KIND, either express or implied. See the License for the
|
||||
# specific language governing permissions and limitations
|
||||
# under the License.
|
||||
import io
|
||||
import tempfile
|
||||
from typing import Any
|
||||
from zipfile import ZipFile
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
from werkzeug.datastructures import FileStorage
|
||||
|
||||
from superset.commands.database.exceptions import DatabaseUploadFailed
|
||||
from superset.commands.database.uploaders.columnar_reader import (
|
||||
ColumnarReader,
|
||||
ColumnarReaderOptions,
|
||||
)
|
||||
from tests.unit_tests.fixtures.common import create_columnar_file
|
||||
|
||||
COLUMNAR_DATA: dict[str, list[Any]] = {
|
||||
"Name": ["name1", "name2", "name3"],
|
||||
"Age": [30, 25, 20],
|
||||
"City": ["city1", "city2", "city3"],
|
||||
"Birth": ["1990-02-01", "1995-02-01", "2000-02-01"],
|
||||
}
|
||||
|
||||
COLUMNAR_WITH_NULLS: dict[str, list[Any]] = {
|
||||
"Name": ["name1", "name2", "name3"],
|
||||
"Age": [None, 25, 20],
|
||||
"City": ["city1", None, "city3"],
|
||||
"Birth": ["1990-02-01", "1995-02-01", "2000-02-01"],
|
||||
}
|
||||
|
||||
|
||||
COLUMNAR_WITH_FLOATS: dict[str, list[Any]] = {
|
||||
"Name": ["name1", "name2", "name3"],
|
||||
"Age": [30.1, 25.1, 20.1],
|
||||
"City": ["city1", "city2", "city3"],
|
||||
"Birth": ["1990-02-01", "1995-02-01", "2000-02-01"],
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"file, options, expected_cols, expected_values",
|
||||
[
|
||||
(
|
||||
create_columnar_file(COLUMNAR_DATA),
|
||||
ColumnarReaderOptions(),
|
||||
["Name", "Age", "City", "Birth"],
|
||||
[
|
||||
["name1", 30, "city1", "1990-02-01"],
|
||||
["name2", 25, "city2", "1995-02-01"],
|
||||
["name3", 20, "city3", "2000-02-01"],
|
||||
],
|
||||
),
|
||||
(
|
||||
create_columnar_file(COLUMNAR_DATA),
|
||||
ColumnarReaderOptions(
|
||||
columns_read=["Name", "Age"],
|
||||
),
|
||||
["Name", "Age"],
|
||||
[
|
||||
["name1", 30],
|
||||
["name2", 25],
|
||||
["name3", 20],
|
||||
],
|
||||
),
|
||||
(
|
||||
create_columnar_file(COLUMNAR_DATA),
|
||||
ColumnarReaderOptions(
|
||||
columns_read=[],
|
||||
),
|
||||
["Name", "Age", "City", "Birth"],
|
||||
[
|
||||
["name1", 30, "city1", "1990-02-01"],
|
||||
["name2", 25, "city2", "1995-02-01"],
|
||||
["name3", 20, "city3", "2000-02-01"],
|
||||
],
|
||||
),
|
||||
(
|
||||
create_columnar_file(COLUMNAR_WITH_NULLS),
|
||||
ColumnarReaderOptions(),
|
||||
["Name", "Age", "City", "Birth"],
|
||||
[
|
||||
["name1", np.nan, "city1", "1990-02-01"],
|
||||
["name2", 25, None, "1995-02-01"],
|
||||
["name3", 20, "city3", "2000-02-01"],
|
||||
],
|
||||
),
|
||||
(
|
||||
create_columnar_file(COLUMNAR_WITH_FLOATS),
|
||||
ColumnarReaderOptions(),
|
||||
["Name", "Age", "City", "Birth"],
|
||||
[
|
||||
["name1", 30.1, "city1", "1990-02-01"],
|
||||
["name2", 25.1, "city2", "1995-02-01"],
|
||||
["name3", 20.1, "city3", "2000-02-01"],
|
||||
],
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_columnar_reader_file_to_dataframe(
|
||||
file, options, expected_cols, expected_values
|
||||
):
|
||||
reader = ColumnarReader(
|
||||
options=options,
|
||||
)
|
||||
df = reader.file_to_dataframe(file)
|
||||
assert df.columns.tolist() == expected_cols
|
||||
actual_values = df.values.tolist()
|
||||
for i in range(len(expected_values)):
|
||||
for j in range(len(expected_values[i])):
|
||||
expected_val = expected_values[i][j]
|
||||
actual_val = actual_values[i][j]
|
||||
|
||||
# Check if both values are NaN
|
||||
if isinstance(expected_val, float) and isinstance(actual_val, float):
|
||||
assert np.isnan(expected_val) == np.isnan(actual_val)
|
||||
else:
|
||||
assert expected_val == actual_val
|
||||
file.close()
|
||||
|
||||
|
||||
def test_excel_reader_wrong_columns_to_read():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(columns_read=["xpto"]),
|
||||
)
|
||||
with pytest.raises(DatabaseUploadFailed) as ex:
|
||||
reader.file_to_dataframe(create_columnar_file(COLUMNAR_DATA))
|
||||
assert (
|
||||
str(ex.value)
|
||||
== (
|
||||
"Parsing error: No match for FieldRef.Name(xpto) in Name: string\n"
|
||||
"Age: int64\n"
|
||||
"City: string\n"
|
||||
"Birth: string\n"
|
||||
"__fragment_index: int32\n"
|
||||
"__batch_index: int32\n"
|
||||
"__last_in_fragment: bool\n"
|
||||
"__filename: string"
|
||||
)
|
||||
!= (
|
||||
"Parsing error: Usecols do not match columns, columns expected but not found: "
|
||||
"['xpto'] (sheet: 0)"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def test_columnar_reader_invalid_file():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(),
|
||||
)
|
||||
with pytest.raises(DatabaseUploadFailed) as ex:
|
||||
reader.file_to_dataframe(FileStorage(io.BytesIO(b"c1"), "test.parquet"))
|
||||
assert str(ex.value) == (
|
||||
"Parsing error: Could not open Parquet input source '<Buffer>': Parquet file "
|
||||
"size is 2 bytes, smaller than the minimum file footer (8 bytes)"
|
||||
)
|
||||
|
||||
|
||||
def test_columnar_reader_zip():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(),
|
||||
)
|
||||
file1 = create_columnar_file(COLUMNAR_DATA, "test1.parquet")
|
||||
file2 = create_columnar_file(COLUMNAR_DATA, "test2.parquet")
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False) as tmp_file1:
|
||||
tmp_file1.write(file1.read())
|
||||
tmp_file1.seek(0)
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False) as tmp_file2:
|
||||
tmp_file2.write(file2.read())
|
||||
tmp_file2.seek(0)
|
||||
|
||||
with tempfile.NamedTemporaryFile(delete=False) as tmp_zip:
|
||||
with ZipFile(tmp_zip, "w") as zip_file:
|
||||
zip_file.write(tmp_file1.name, "test1.parquet")
|
||||
zip_file.write(tmp_file2.name, "test2.parquet")
|
||||
tmp_zip.seek(0) # Reset file pointer to beginning
|
||||
df = reader.file_to_dataframe(FileStorage(tmp_zip, "test.zip"))
|
||||
assert df.columns.tolist() == ["Name", "Age", "City", "Birth"]
|
||||
assert df.values.tolist() == [
|
||||
["name1", 30, "city1", "1990-02-01"],
|
||||
["name2", 25, "city2", "1995-02-01"],
|
||||
["name3", 20, "city3", "2000-02-01"],
|
||||
["name1", 30, "city1", "1990-02-01"],
|
||||
["name2", 25, "city2", "1995-02-01"],
|
||||
["name3", 20, "city3", "2000-02-01"],
|
||||
]
|
||||
|
||||
|
||||
def test_columnar_reader_bad_parquet_in_zip():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(),
|
||||
)
|
||||
with tempfile.NamedTemporaryFile(delete=False) as tmp_zip:
|
||||
with ZipFile(tmp_zip, "w") as zip_file:
|
||||
zip_file.writestr("test1.parquet", b"bad parquet file")
|
||||
zip_file.writestr("test2.parquet", b"bad parquet file")
|
||||
tmp_zip.seek(0) # Reset file pointer to beginning
|
||||
with pytest.raises(DatabaseUploadFailed) as ex:
|
||||
reader.file_to_dataframe(FileStorage(tmp_zip, "test.zip"))
|
||||
assert str(ex.value) == (
|
||||
"Parsing error: Could not open Parquet input source '<Buffer>': "
|
||||
"Parquet magic bytes not found in footer. "
|
||||
"Either the file is corrupted or this is not a parquet file."
|
||||
)
|
||||
|
||||
|
||||
def test_columnar_reader_bad_zip():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(),
|
||||
)
|
||||
with pytest.raises(DatabaseUploadFailed) as ex:
|
||||
reader.file_to_dataframe(FileStorage(io.BytesIO(b"bad zip file"), "test.zip"))
|
||||
assert str(ex.value) == "Not a valid ZIP file"
|
||||
|
||||
|
||||
def test_columnar_reader_metadata():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(),
|
||||
)
|
||||
file = create_columnar_file(COLUMNAR_DATA)
|
||||
metadata = reader.file_metadata(file)
|
||||
column_names = sorted(metadata["items"][0]["column_names"])
|
||||
assert column_names == ["Age", "Birth", "City", "Name"]
|
||||
assert metadata["items"][0]["sheet_name"] is None
|
||||
|
||||
|
||||
def test_columnar_reader_metadata_invalid_file():
|
||||
reader = ColumnarReader(
|
||||
options=ColumnarReaderOptions(),
|
||||
)
|
||||
with pytest.raises(DatabaseUploadFailed) as ex:
|
||||
reader.file_metadata(FileStorage(io.BytesIO(b"c1"), "test.parquet"))
|
||||
assert str(ex.value) == (
|
||||
"Parsing error: Parquet file size is 2 bytes, "
|
||||
"smaller than the minimum file footer (8 bytes)"
|
||||
)
|
||||
Reference in New Issue
Block a user