forked from lance-format/lance
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathtest_scan.py
More file actions
148 lines (118 loc) · 4.25 KB
/
test_scan.py
File metadata and controls
148 lines (118 loc) · 4.25 KB
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
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The Lance Authors
import random
from pathlib import Path
import lance
import pyarrow as pa
import pyarrow.compute as pc
import pytest
NUM_ROWS = 10_000
@pytest.mark.parametrize(
"array_factory",
[
lambda: pa.array(range(NUM_ROWS), type=pa.int32()),
lambda: pc.random(NUM_ROWS),
lambda: pa.array(
[random.choice(["hello", "world", "today"]) for _ in range(NUM_ROWS)],
type=pa.string(),
),
lambda: pa.array(
[random.choice(["hello", "world", "today"]) for _ in range(NUM_ROWS)],
type=pa.dictionary(pa.int8(), pa.string()),
),
lambda: pa.FixedSizeListArray.from_arrays(
pc.random(NUM_ROWS * 128).cast(pa.float32()), 128
),
],
ids=["i32", "f64", "string", "dictionary", "vector"],
)
@pytest.mark.benchmark(group="scan_single_column")
def test_scan_integer(tmp_path: Path, benchmark, array_factory):
values = array_factory()
table = pa.table({"values": values})
dataset = lance.write_dataset(table, tmp_path)
result = benchmark(
dataset.to_table,
)
assert result.num_rows == NUM_ROWS
@pytest.fixture(scope="module")
def sample_dataset(tmpdir_factory):
tmp_path = Path(tmpdir_factory.mktemp("data"))
table = pa.table(
{
"i": pa.array(range(NUM_ROWS), type=pa.int32()),
"f": pc.random(NUM_ROWS).cast(pa.float32()),
"s": pa.array(
[random.choice(["hello", "world", "today"]) for _ in range(NUM_ROWS)],
type=pa.string(),
),
"fsl": pa.FixedSizeListArray.from_arrays(
pc.random(NUM_ROWS * 128).cast(pa.float32()), 128
),
"blob": pa.array(
[
random.choice(
[
random.randbytes(100 * 1024),
random.randbytes(100 * 1024),
random.randbytes(100 * 1024),
]
)
for _ in range(NUM_ROWS)
],
type=pa.binary(),
),
}
)
return lance.write_dataset(table, tmp_path)
@pytest.mark.benchmark(group="scan_table")
def test_scan_table_full(benchmark, sample_dataset):
result = benchmark(
sample_dataset.to_table,
)
assert result.num_rows == NUM_ROWS
@pytest.mark.benchmark(group="scan_table")
def test_scan_table_project(benchmark, sample_dataset):
result = benchmark(sample_dataset.to_table, columns=["i", "f"])
assert result.schema.names == ["i", "f"]
assert result.num_rows == NUM_ROWS
@pytest.mark.parametrize("keep_percent", [0.1, 0.5, 0.9])
@pytest.mark.benchmark(group="scan_table")
def test_scan_table_filter_project(benchmark, sample_dataset, keep_percent):
result = benchmark(
sample_dataset.to_table,
filter=f"f < {keep_percent}",
columns=["i", "blob"],
)
assert result.schema.names == ["i", "blob"]
@pytest.mark.parametrize("keep_percent", [0.1, 0.5, 0.9])
@pytest.mark.benchmark(group="scan_table")
def test_scan_table_filter_full(benchmark, sample_dataset, keep_percent):
result = benchmark(
sample_dataset.to_table,
filter=f"f < {keep_percent}",
)
assert result.schema.names == ["i", "f", "s", "fsl", "blob"]
@pytest.mark.benchmark(group="filter_table")
def test_filter_for_range(benchmark, sample_dataset):
result = benchmark(
sample_dataset.to_table,
filter="i > 1000 and i < 5000",
)
assert result.schema.names == ["i", "f", "s", "fsl", "blob"]
@pytest.mark.benchmark(group="filter_table")
def test_filter_for_row(benchmark, sample_dataset):
result = benchmark(
sample_dataset.to_table,
filter="i = 4200",
)
assert result.num_rows == 1
assert result.schema.names == ["i", "f", "s", "fsl", "blob"]
@pytest.mark.benchmark(group="filter_table")
def test_filter_multiple(benchmark, sample_dataset):
result = benchmark(
sample_dataset.to_table,
filter="i > 1000 and i < 5000 and s in ('hello', 'world')",
)
assert result.num_rows > 1
assert result.schema.names == ["i", "f", "s", "fsl", "blob"]