mirror of
https://ghfast.top/https://github.com/aeroxw/tick-stock-panel.git
synced 2026-09-12 14:24:15 +08:00
82 lines
2.6 KiB
Python
82 lines
2.6 KiB
Python
from datetime import date
|
|
|
|
import polars as pl
|
|
|
|
from app.parquet import scan_daily_parquet, scan_enriched_parquet
|
|
|
|
|
|
def test_partitioned_daily_scan_tolerates_added_quote_ts(tmp_path):
|
|
old_part = tmp_path / "kline_daily" / "date=2026-07-08" / "part.parquet"
|
|
new_part = tmp_path / "kline_daily" / "date=2026-07-09" / "part.parquet"
|
|
old_part.parent.mkdir(parents=True)
|
|
new_part.parent.mkdir(parents=True)
|
|
|
|
pl.DataFrame({
|
|
"symbol": ["600000.SH"],
|
|
"date": [date(2026, 7, 8)],
|
|
"open": [10.0],
|
|
"high": [10.5],
|
|
"low": [9.8],
|
|
"close": [10.2],
|
|
"volume": [1000.0],
|
|
"amount": [10200.0],
|
|
}).write_parquet(old_part)
|
|
|
|
pl.DataFrame({
|
|
"symbol": ["600000.SH"],
|
|
"date": [date(2026, 7, 9)],
|
|
"open": [10.2],
|
|
"high": [10.8],
|
|
"low": [10.1],
|
|
"close": [10.6],
|
|
"volume": [1200],
|
|
"amount": [12720.0],
|
|
"quote_ts": [1783560600000],
|
|
}).write_parquet(new_part)
|
|
|
|
df = scan_daily_parquet(str(tmp_path / "kline_daily" / "**" / "*.parquet")).sort("date").collect()
|
|
|
|
assert df.height == 2
|
|
assert df.schema["volume"] == pl.Float64
|
|
assert df.schema["quote_ts"] == pl.Int64
|
|
assert df["quote_ts"].to_list() == [None, 1783560600000]
|
|
|
|
|
|
def test_partitioned_enriched_scan_tolerates_added_quote_ts(tmp_path):
|
|
base = tmp_path / "kline_daily_enriched"
|
|
old_part = base / "date=2026-07-08" / "part.parquet"
|
|
new_part = base / "date=2026-07-09" / "part.parquet"
|
|
old_part.parent.mkdir(parents=True)
|
|
new_part.parent.mkdir(parents=True)
|
|
|
|
common_old = {
|
|
"symbol": ["600000.SH"],
|
|
"date": [date(2026, 7, 8)],
|
|
"open": [10.0],
|
|
"high": [10.5],
|
|
"low": [9.8],
|
|
"close": [10.2],
|
|
"volume": [1000.0],
|
|
"amount": [10200.0],
|
|
"raw_close": [10.2],
|
|
"raw_high": [10.5],
|
|
"raw_low": [9.8],
|
|
"turnover_rate": [1.1],
|
|
"consecutive_limit_ups": pl.Series([0], dtype=pl.UInt32),
|
|
"consecutive_limit_downs": pl.Series([0], dtype=pl.UInt32),
|
|
}
|
|
pl.DataFrame(common_old).write_parquet(old_part)
|
|
|
|
common_new = dict(common_old)
|
|
common_new["date"] = [date(2026, 7, 9)]
|
|
common_new["volume"] = [1200]
|
|
common_new["quote_ts"] = [1783560600000]
|
|
pl.DataFrame(common_new).write_parquet(new_part)
|
|
|
|
df = scan_enriched_parquet(str(base / "**" / "*.parquet")).sort("date").collect()
|
|
|
|
assert df.height == 2
|
|
assert df.schema["volume"] == pl.Float64
|
|
assert df.schema["quote_ts"] == pl.Int64
|
|
assert df["quote_ts"].to_list() == [None, 1783560600000]
|